From 46d41431ba88c7b741a5aa78bb8483380bf979b2 Mon Sep 17 00:00:00 2001 From: Miles Olson Date: Mon, 16 Mar 2026 10:31:57 -0700 Subject: [PATCH] Refactor Objective, OutcomeConstraint to be expression-based; remove Metrics from OptimizationConfig (#5000) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/5000 This diff is a major refactor. In this diff we remove Metric instances from both Objective and OutcomeConstraint (and transitively, also off OptimizationConfig) and instead add them directly to the Experiment. In their place, we use the "expression" syntax and Sympy parsing from the Ax API directly on these classes, allowing users to specify `Objective(expression="ne1 + 2*ne2")` or `OutcomeConstraint(expression="qps >= 7000")` directly. This has a number of benefits: * Dependency spread is contained. OptimizationConfig is now able to be loaded without loading the any specific Metric implementation * Major separation of concerns: OptimizationConfig is now solely concerned with the goals of the optimization and not also implicitly concerned with fetching logic * Bringing Expressions directly onto Objective and OutcomeConstraint allows us to deprecated ScalarizedObjective, MultiObjective, ScalarizedOutcomeConstraint, and ObjectiveThreshold * Massive simplification for eventual upstreaming of MultiTypeExperiment functionality into base Experiment This is also a necessary step in our storage redesign in order to make metrics "live" only in one place. See this doc for details: https://docs.google.com/document/d/1I8FSQJ05_WHXFHBtNaL3dTAdAbvKfADwMRlIdq6BEV8/edit?tab=t.0#heading=h.s4r4rlrnva07 The refactor touches several core modules, please focus review on these: * ax/core/objective.py * ax/core/outcome_constraint.py * ax/core/optimization_config.py * ax/core/experiment.py These changes mandate many many updates to callsites which is ballooning the size of the diff. Changes are summarized here. # Core changes: * Objective, OutcomeConstraint classes now takes an expression string (e.g., "accuracy", "-loss", "2*acc + recall", "acc, -loss") * Old-style constructors still work, but they discard the Metric after initialization * Objective.metric property no longer exists. If a user wants to get the specific metric they retrieve its name via get_names and pass that string into Experiment.get_metric(name: str) -> Metric * isinstance checks for MultiObective, ScalarizedObjective, etc are replaced with is_multi_objective and is_scalarized_objective properties * Experiment.tracking_metrics is a property not an attr, now experiment simply has a collection of Metrics called "metrics" and tracking_metrics is a property which returns the metrics not on the optimization config * **Parsing is directly lifted and shifted from ax/api** * **Both JSON and SQL encoders/decoders carefully reconstruct the pre-refactor structure, allowing existing storage to be unaffected** # Other changes ## **ax/adapter/** - **adapter_utils.py** - Functions updated to accept experiment metrics for lookup. - All extraction functions use metric names. - **base.py, torch.py, transforms/** - Updated to thread experiment metrics through extraction functions. - Metric access patterns updated. ## **ax/analysis/** - Updated all metric access patterns to use metric names, not Metric objects. ## **ax/api/** - **client.py** - Simplified metric overwriting logic: metrics only live on Experiment. - **utils/instantiation/from_string.py** - Simplified parsing: constructs Objective directly from expression string. ## **ax/benchmark/** - Updated objective construction to use expression-based interface. ## **ax/core/** - **experiment.py** - Experiment now maintains a flat list of metrics (`_metrics: dict[str, Metric]`). - No distinction between tracking and optimization metrics. - Objective and constraints reference metrics by name, not object. - **objectives.py** - Unified Objective class: now takes an expression string, supporting single, scalarized, and multi-objective optimization. - Deprecated MultiObjective and ScalarizedObjective classes; their functionality is merged into Objective. - All metric access is via metric names/weights parsed from the expression, not direct Metric objects. - **optimization_config.py** - Unified OptimizationConfig class: handles both single and multi-objective cases. - Removed MultiObjectiveOptimizationConfig. - All config logic uses expression-based Objective and OutcomeConstraint. - **outcome_constraint.py** - Unified OutcomeConstraint class: now takes an expression string, supporting scalarized constraints. - Removed ObjectiveThreshold and ScalarizedOutcomeConstraint classes. - All validation and access is via metric names/weights parsed from the expression. - **trial.py** - Updated to reference metrics via `objective.metric_names[0]` instead of `objective.metric.name`. - **utils.py** - Updated metric access patterns to use metric names. ## **ax/early_stopping/** - Updated to reference metric signatures via experiment metrics, not via Objective. ## **ax/plot/** - Updated metric access patterns to use metric names. ## **ax/service/** - **ax_client.py** - Updated to reference metrics via `objective.metric_names[0]`. - Simplified metric overwriting logic. - **utils/best_point.py** - Updated metric access patterns to use metric names. ## **ax/storage/json_store/** - Encoder: serializes objectives as expression strings. - Decoder: converts legacy formats (with embedded Metric dicts) to expression-based objects. ## **ax/storage/sqa_store/** - Encoder: stores expression string in JSON column. - Decoder: reconstructs expression strings from MetricIntent-based SQA rows. ## **Tests** - All test constructors rewritten to use expression strings. - Assertions updated for new property names. - Added tests for expression parsing, multi-objective, scalarized, and minimize detection. Reviewed By: lena-kashtelyan Differential Revision: D93520819 --- ax/adapter/adapter_utils.py | 97 ++-- ax/adapter/base.py | 12 +- ax/adapter/cross_validation.py | 10 +- ax/adapter/discrete.py | 9 +- ax/adapter/tests/test_adapter_utils.py | 21 +- ax/adapter/tests/test_base_adapter.py | 29 +- ax/adapter/tests/test_cross_validation.py | 7 +- ax/adapter/tests/test_discrete_adapter.py | 12 +- .../tests/test_hierarchical_search_space.py | 6 +- ax/adapter/tests/test_torch_adapter.py | 13 +- ax/adapter/tests/test_torch_moo_adapter.py | 69 ++- ax/adapter/tests/test_transform_utils.py | 2 +- ax/adapter/tests/test_utils.py | 29 +- ax/adapter/torch.py | 44 +- ax/adapter/transforms/base.py | 16 + ax/adapter/transforms/bilog_y.py | 3 +- ax/adapter/transforms/derelativize.py | 117 ++-- ax/adapter/transforms/log_y.py | 77 ++- .../transforms/objective_as_constraint.py | 57 +- ax/adapter/transforms/power_transform_y.py | 134 +++-- ax/adapter/transforms/relativize.py | 48 +- ax/adapter/transforms/standardize_y.py | 161 ++++-- .../transforms/stratified_standardize_y.py | 167 ++++-- .../tests/test_add_execution_viability.py | 2 +- .../transforms/tests/test_base_transform.py | 6 +- .../tests/test_derelativize_transform.py | 66 ++- .../transforms/tests/test_log_y_transform.py | 15 +- .../tests/test_objective_as_constraint.py | 44 +- .../tests/test_power_y_transform.py | 7 +- .../tests/test_relativize_transform.py | 21 +- .../tests/test_standardize_y_transform.py | 62 ++- ...test_stratified_standardize_y_transform.py | 115 ++-- .../tests/test_winsorize_transform.py | 55 +- ax/adapter/transforms/winsorize.py | 34 +- ax/analysis/diagnostics.py | 2 +- .../healthcheck/constraints_feasibility.py | 6 +- .../healthcheck/early_stopping_healthcheck.py | 6 +- .../tests/test_complexity_rating.py | 7 + ax/analysis/insights.py | 2 +- ax/analysis/plotly/constraint_feasibility.py | 10 +- .../plotly/objective_p_feasible_frontier.py | 4 +- ax/analysis/plotly/p_feasible.py | 6 +- .../plotly/tests/test_marginal_effects.py | 2 +- .../test_objective_p_feasible_frontier.py | 10 +- ax/analysis/plotly/utility_progression.py | 14 +- ax/analysis/plotly/utils.py | 6 +- ax/analysis/results.py | 5 +- ax/analysis/tests/test_metric_summary.py | 2 +- ax/analysis/tests/test_results.py | 7 +- ax/analysis/trials.py | 2 +- ax/analysis/utils.py | 12 +- ax/api/client.py | 90 ++- ax/api/tests/test_client.py | 68 ++- ax/api/utils/instantiation/from_string.py | 180 +----- .../instantiation/tests/test_from_string.py | 171 +----- ax/benchmark/benchmark.py | 14 +- ax/benchmark/benchmark_problem.py | 163 +++--- .../surrogate/lcbench/early_stopping.py | 7 +- .../problems/surrogate/lcbench/utils.py | 3 +- ax/benchmark/problems/synthetic/bandit.py | 3 +- .../synthetic/discretized/mixed_integer.py | 3 +- .../problems/synthetic/from_botorch.py | 5 +- .../problems/synthetic/hss/jenatton.py | 3 +- ax/benchmark/testing/benchmark_stubs.py | 11 +- .../surrogate/hss/test_cifar10_surrogate.py | 3 +- .../hss/test_fashion_mnist_surrogate.py | 3 +- .../surrogate/hss/test_mnist_surrogate.py | 4 +- .../surrogate/lcbench/test_early_stopping.py | 2 +- .../problems/synthetic/hss/test_jenatton.py | 12 +- .../problems/synthetic/test_from_botorch.py | 54 +- ax/benchmark/tests/test_benchmark.py | 61 +- ax/benchmark/tests/test_benchmark_problem.py | 91 ++- ax/core/experiment.py | 253 ++++++--- ax/core/multi_type_experiment.py | 22 +- ax/core/objective.py | 356 +++++++----- ax/core/optimization_config.py | 184 +++--- ax/core/outcome_constraint.py | 527 +++++++++--------- ax/core/tests/test_derived_metric.py | 8 +- ax/core/tests/test_experiment.py | 122 +++- ax/core/tests/test_multi_type_experiment.py | 4 +- ax/core/tests/test_objective.py | 341 +++++++++--- ax/core/tests/test_optimization_config.py | 461 ++++++++------- ax/core/tests/test_outcome_constraint.py | 506 ++++++++++------- ax/core/tests/test_trial.py | 8 +- ax/core/tests/test_utils.py | 12 +- ax/core/trial.py | 2 +- ax/core/utils.py | 18 +- ax/early_stopping/dispatch.py | 5 +- ax/early_stopping/experiment_replay.py | 1 + ax/early_stopping/strategies/base.py | 21 +- .../tests/test_experiment_replay.py | 8 +- ax/early_stopping/tests/test_strategies.py | 20 +- .../tests/test_dispatch_utils.py | 8 +- .../tests/test_transition_criterion.py | 11 + ax/global_stopping/strategies/improvement.py | 10 +- ax/global_stopping/tests/test_strategies.py | 14 +- ax/metrics/tests/test_map_replay.py | 2 + ax/orchestration/orchestrator.py | 4 +- ax/orchestration/tests/test_orchestrator.py | 27 +- ax/plot/helper.py | 2 +- ax/plot/pareto_frontier.py | 21 +- ax/plot/pareto_utils.py | 15 +- ax/plot/tests/test_pareto_utils.py | 19 +- ax/runners/tests/test_map_replay.py | 1 + ax/runners/tests/test_torchx.py | 9 +- ax/service/ax_client.py | 47 +- ax/service/tests/test_ax_client.py | 159 +++--- ax/service/tests/test_best_point.py | 49 +- ax/service/tests/test_best_point_utils.py | 135 +++-- ax/service/tests/test_instantiation_utils.py | 27 +- ax/service/tests/test_report_utils.py | 50 +- ax/service/utils/best_point.py | 138 +++-- ax/service/utils/best_point_mixin.py | 8 +- ax/service/utils/instantiation.py | 47 +- ax/service/utils/report_utils.py | 54 +- ax/storage/json_store/decoder.py | 15 +- ax/storage/json_store/encoders.py | 171 +++++- .../json_store/tests/test_json_store.py | 18 +- ax/storage/sqa_store/decoder.py | 165 ++++-- ax/storage/sqa_store/encoder.py | 246 +++++--- ax/storage/sqa_store/tests/test_sqa_store.py | 48 +- ax/storage/sqa_store/utils.py | 5 +- ax/utils/common/complexity_utils.py | 9 +- ax/utils/common/string_utils.py | 56 +- ax/utils/common/sympy.py | 157 +++++- ax/utils/common/tests/test_string_utils.py | 7 +- ax/utils/testing/core_stubs.py | 188 +++++-- ax/utils/testing/utils.py | 2 +- .../external_generation_node.ipynb | 2 +- 129 files changed, 4646 insertions(+), 2800 deletions(-) diff --git a/ax/adapter/adapter_utils.py b/ax/adapter/adapter_utils.py index c42889dca82..33ab609fec5 100644 --- a/ax/adapter/adapter_utils.py +++ b/ax/adapter/adapter_utils.py @@ -23,7 +23,7 @@ ) from ax.core.arm import Arm from ax.core.experiment import Experiment -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, @@ -206,6 +206,7 @@ def extract_objective_thresholds( objective_thresholds: TRefPoint, objective: Objective, outcomes: list[str], + experiment: Experiment, ) -> npt.NDArray | None: """Extracts objective thresholds' values, in the order of `outcomes`. @@ -221,6 +222,7 @@ def extract_objective_thresholds( objective_thresholds: Objective thresholds to extract values from. objective: The corresponding Objective, for validation purposes. outcomes: n-length list of names of metrics. + experiment: The experiment, used to map metric names to signatures. Returns: (n,) array of thresholds @@ -230,15 +232,19 @@ def extract_objective_thresholds( objective_threshold_dict = {} for ot in objective_thresholds: + ot_signature = experiment.get_metric(ot.metric_names[0]).signature if ot.relative: raise ValueError( - f"Objective {ot.metric.signature} has a relative threshold that " + f"Objective {ot_signature} has a relative threshold that " f"is not supported here." ) - objective_threshold_dict[ot.metric.signature] = ot.bound + objective_threshold_dict[ot_signature] = ot.bound # Check that all thresholds correspond to a metric. - if set(objective_threshold_dict.keys()).difference(set(objective.metric_names)): + obj_metric_signatures = [ + experiment.get_metric(name).signature for name in objective.metric_names + ] + if set(objective_threshold_dict.keys()).difference(set(obj_metric_signatures)): raise ValueError( "Some objective thresholds do not have corresponding metrics. " f"Got {objective_thresholds=} and {objective=}." @@ -252,7 +258,9 @@ def extract_objective_thresholds( return obj_t -def extract_objective_weights(objective: Objective, outcomes: list[str]) -> npt.NDArray: +def extract_objective_weights( + objective: Objective, outcomes: list[str], experiment: Experiment +) -> npt.NDArray: """Extract a weights for objectives. Weights are for a maximization problem. @@ -268,29 +276,24 @@ def extract_objective_weights(objective: Objective, outcomes: list[str]) -> npt. Args: objective: Objective to extract weights from. - outcomes: n-length list of names of metrics. + outcomes: n-length list of metric signatures. + experiment: The experiment, used to map metric names to signatures. Returns: n-length array of weights. """ objective_weights = np.zeros(len(outcomes)) - if isinstance(objective, ScalarizedObjective): - s = -1.0 if objective.minimize else 1.0 - for obj_metric, obj_weight in objective.metric_weights: - objective_weights[outcomes.index(obj_metric.signature)] = obj_weight * s - elif isinstance(objective, MultiObjective): - for obj in objective.objectives: - s = -1.0 if obj.minimize else 1.0 - objective_weights[outcomes.index(obj.metric.signature)] = s - else: - s = -1.0 if objective.minimize else 1.0 - objective_weights[outcomes.index(objective.metric.signature)] = s + # metric_weights returns sign-encoded (name, weight) tuples for all + # objective types (single, scalarized, multi). + for obj_metric_name, obj_weight in objective.metric_weights: + sig = experiment.get_metric(obj_metric_name).signature + objective_weights[outcomes.index(sig)] = obj_weight return objective_weights def extract_objective_weight_matrix( - objective: Objective, outcomes: list[str] + objective: Objective, outcomes: list[str], experiment: Experiment ) -> npt.NDArray: """Extract a 2D weight matrix for objectives. @@ -304,23 +307,31 @@ def extract_objective_weight_matrix( Args: objective: Objective to extract weights from. - outcomes: n-length list of names of metrics. + outcomes: n-length list of signatures of metrics. Returns: ``(n_objectives, n)`` array of weights. """ - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: rows: list[npt.NDArray] = [] - for obj in objective.objectives: - rows.append(extract_objective_weights(obj, outcomes)) + for name, weight in objective.metric_weights: + rows.append( + extract_objective_weights( + objective=Objective(expression=f"{weight} * {name}"), + outcomes=outcomes, + experiment=experiment, + ) + ) return np.stack(rows, axis=0) else: # Single row – covers Objective and ScalarizedObjective - return extract_objective_weights(objective, outcomes).reshape(1, -1) + return extract_objective_weights(objective, outcomes, experiment).reshape(1, -1) def extract_outcome_constraints( - outcome_constraints: list[OutcomeConstraint], outcomes: list[str] + outcome_constraints: list[OutcomeConstraint], + outcomes: list[str], + experiment: Experiment, ) -> TBounds: if len(outcome_constraints) == 0: return None @@ -330,11 +341,11 @@ def extract_outcome_constraints( for i, c in enumerate(outcome_constraints): s = 1 if c.op == ComparisonOp.LEQ else -1 if isinstance(c, ScalarizedOutcomeConstraint): - for c_metric, c_weight in c.metric_weights: - j = outcomes.index(c_metric.signature) + for c_metric_name, c_weight in c.metric_weights: + j = outcomes.index(experiment.get_metric(c_metric_name).signature) A[i, j] = s * c_weight else: - j = outcomes.index(c.metric.signature) + j = outcomes.index(experiment.get_metric(c.metric_names[0]).signature) A[i, j] = s b[i, 0] = s * c.bound return (A, b) @@ -645,16 +656,20 @@ def get_pareto_frontier_and_configs( ) # Extract weights, constraints, and objective_thresholds objective_weights = extract_objective_weight_matrix( - objective=optimization_config.objective, outcomes=adapter.outcomes + objective=optimization_config.objective, + outcomes=adapter.outcomes, + experiment=adapter._experiment, ) outcome_constraints = extract_outcome_constraints( outcome_constraints=optimization_config.outcome_constraints, outcomes=adapter.outcomes, + experiment=adapter._experiment, ) obj_t = extract_objective_thresholds( objective_thresholds=optimization_config.objective_thresholds, objective=optimization_config.objective, outcomes=adapter.outcomes, + experiment=adapter._experiment, ) if obj_t is not None: obj_t = array_to_tensor(obj_t) @@ -1113,6 +1128,7 @@ def observation_features_to_array( def feasible_hypervolume( optimization_config: MultiObjectiveOptimizationConfig, values: dict[str, npt.NDArray], + experiment: Experiment, ) -> npt.NDArray: """Compute the feasible hypervolume each iteration. @@ -1121,34 +1137,35 @@ def feasible_hypervolume( values: Dictionary from metric name to array of value at each iteration (each array is `n`-dim). If optimization config contains outcome constraints, values for them must be present in `values`. + experiment: The experiment, used to map metric names to signatures. Returns: Array of feasible hypervolumes. """ # Get objective at each iteration obj_threshold_dict = { - ot.metric.signature: ot.bound for ot in optimization_config.objective_thresholds + experiment.get_metric(ot.metric_names[0]).signature: ot.bound + for ot in optimization_config.objective_thresholds } - f_vals = np.hstack( - [ - values[m.signature].reshape(-1, 1) - for m in optimization_config.objective.metrics - ] - ) - obj_thresholds = np.array( - [obj_threshold_dict[m.signature] for m in optimization_config.objective.metrics] - ) + obj_metric_names = optimization_config.objective.metric_names + obj_metrics = [experiment.get_metric(name) for name in obj_metric_names] + f_vals = np.hstack([values[m.signature].reshape(-1, 1) for m in obj_metrics]) + obj_thresholds = np.array([obj_threshold_dict[m.signature] for m in obj_metrics]) # Set infeasible points to be the objective threshold for oc in optimization_config.outcome_constraints: if oc.relative: raise ValueError( "Benchmark aggregation does not support relative constraints" ) - g = values[oc.metric.signature] + oc_sig = experiment.get_metric(oc.metric_names[0]).signature + g = values[oc_sig] feas = g <= oc.bound if oc.op == ComparisonOp.LEQ else g >= oc.bound f_vals[~feas] = obj_thresholds + # Derive objective directions from the objective's metric_weights. + # Positive weight = maximize, negative weight = minimize. + obj_weight_dict = dict(optimization_config.objective.metric_weights) obj_weights = np.array( - [-1 if m.lower_is_better else 1 for m in optimization_config.objective.metrics] + [1 if obj_weight_dict[name] > 0 else -1 for name in obj_metric_names] ) obj_thresholds = obj_thresholds * obj_weights f_vals = f_vals * obj_weights diff --git a/ax/adapter/base.py b/ax/adapter/base.py index 74e02d1e487..338522eb216 100644 --- a/ax/adapter/base.py +++ b/ax/adapter/base.py @@ -201,7 +201,7 @@ def __init__( "Optimization config is required when " "`fit_tracking_metrics` is False." ) - self.outcomes = sorted(self._optimization_config.metrics.keys()) + self.outcomes = sorted(self._optimization_config.metric_names) # Set training data (in the raw / untransformed space). This also omits # out-of-design and abandoned observations depending on the corresponding flags. @@ -466,10 +466,14 @@ def _set_status_quo(self, experiment: Experiment) -> None: ) return - if has_map_metrics(optimization_config=self._optimization_config): + if has_map_metrics( + metrics=experiment.get_metrics( + metric_names=[*self._optimization_config.metric_names] + ) + ): self._status_quo = _combine_multiple_status_quo_observations( status_quo_observations=status_quo_observations, - metrics=set(none_throws(self._optimization_config).metrics), + metrics=none_throws(self._optimization_config).metric_names, ) else: logger.warning( @@ -689,7 +693,7 @@ def _get_transformed_gen_args( # Check that the optimization config has the same metrics as # the original one. Otherwise, we may attempt to optimize over # metrics that do not have a fitted model. - outcomes = set(optimization_config.metrics.keys()) + outcomes = optimization_config.metric_names if not outcomes.issubset(self.outcomes): raise UnsupportedError( "When fit_tracking_metrics is False, the optimization config " diff --git a/ax/adapter/cross_validation.py b/ax/adapter/cross_validation.py index e638367ffc1..a7fb113c07d 100644 --- a/ax/adapter/cross_validation.py +++ b/ax/adapter/cross_validation.py @@ -21,6 +21,7 @@ from ax.adapter.data_utils import ExperimentData from ax.adapter.observation_utils import unwrap_observation_data from ax.adapter.torch import TorchAdapter +from ax.core.experiment import Experiment from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.optimization_config import OptimizationConfig from ax.exceptions.core import UnsupportedError @@ -573,6 +574,7 @@ def assess_model_fit( def has_good_opt_config_model_fit( optimization_config: OptimizationConfig, assess_model_fit_result: AssessModelFitResult, + experiment: Experiment, ) -> bool: """Assess model fit for given diagnostics results across the optimization config metrics @@ -584,7 +586,8 @@ def has_good_opt_config_model_fit( Args: optimization_config: Objective/Outcome constraint metrics to assess - diagnostics: Output of compute_diagnostics + assess_model_fit_result: Output of assess_model_fit + experiment: The experiment, used to map metric names to signatures. Returns: Two dictionaries, one for good metrics, one for bad metrics, each @@ -594,8 +597,9 @@ def has_good_opt_config_model_fit( # Bad fit criteria: Any objective metrics are poorly fit # TODO[]: Incl. outcome constraints in assessment has_good_opt_config_fit = all( - (m.signature in assess_model_fit_result.good_fit_metrics_to_fisher_score) - for m in optimization_config.objective.metrics + experiment.get_metric(name).signature + in assess_model_fit_result.good_fit_metrics_to_fisher_score + for name in optimization_config.objective.metric_names ) return has_good_opt_config_fit diff --git a/ax/adapter/discrete.py b/ax/adapter/discrete.py index ad2c7b1414d..7645427b97e 100644 --- a/ax/adapter/discrete.py +++ b/ax/adapter/discrete.py @@ -163,13 +163,18 @@ def _gen( objective_weights = None outcome_constraints = None else: - validate_transformed_optimization_config(optimization_config, self.outcomes) + validate_transformed_optimization_config( + optimization_config, self.outcomes, experiment=self._experiment + ) objective_weights = extract_objective_weights( - objective=optimization_config.objective, outcomes=self.outcomes + objective=optimization_config.objective, + outcomes=self.outcomes, + experiment=self._experiment, ) outcome_constraints = extract_outcome_constraints( outcome_constraints=optimization_config.outcome_constraints, outcomes=self.outcomes, + experiment=self._experiment, ) # Get fixed features diff --git a/ax/adapter/tests/test_adapter_utils.py b/ax/adapter/tests/test_adapter_utils.py index 4b618db2981..7f9bd9899af 100644 --- a/ax/adapter/tests/test_adapter_utils.py +++ b/ax/adapter/tests/test_adapter_utils.py @@ -43,6 +43,7 @@ from ax.utils.common.hash_utils import compute_lilo_input_hash from ax.utils.common.testutils import TestCase from ax.utils.testing.core_stubs import ( + get_branin_experiment, get_experiment_with_observations, get_hierarchical_search_space, get_search_space_for_range_values, @@ -61,7 +62,7 @@ def test_feasible_hypervolume(self) -> None: ), outcome_constraints=[ OutcomeConstraint( - mc, + metric=mc, op=ComparisonOp.GEQ, bound=0, relative=False, @@ -78,6 +79,11 @@ def test_feasible_hypervolume(self) -> None: ), ], ) + experiment = Experiment( + search_space=SearchSpace(parameters=[]), + optimization_config=optimization_config, + tracking_metrics=[mc], + ) feas_hv = feasible_hypervolume( optimization_config, values={ @@ -106,6 +112,7 @@ def test_feasible_hypervolume(self) -> None: ] ), }, + experiment=experiment, ) self.assertEqual(list(feas_hv), [0.0, 0.0, 1.0, 1.0]) @@ -537,20 +544,24 @@ def test_can_map_to_binary(self) -> None: def test_extract_objective_weight_matrix(self) -> None: m1, m2, m3 = Metric(name="m1"), Metric(name="m2"), Metric(name="m3") outcomes = ["m1", "m2", "m3"] + experiment = get_branin_experiment() + experiment.add_metric(m1) + experiment.add_metric(m2) + experiment.add_metric(m3) # Single Objective: one row, nonzero only in matching column. obj = Objective(metric=m1, minimize=False) - result = extract_objective_weight_matrix(obj, outcomes) + result = extract_objective_weight_matrix(obj, outcomes, experiment) np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0]]) # Minimization flips the sign. obj_min = Objective(metric=m2, minimize=True) - result = extract_objective_weight_matrix(obj_min, outcomes) + result = extract_objective_weight_matrix(obj_min, outcomes, experiment) np.testing.assert_array_equal(result, [[0.0, -1.0, 0.0]]) # ScalarizedObjective: single row with multiple nonzero entries. scal = ScalarizedObjective(metrics=[m1, m3], weights=[0.3, 0.7], minimize=False) - result = extract_objective_weight_matrix(scal, outcomes) + result = extract_objective_weight_matrix(scal, outcomes, experiment) np.testing.assert_array_almost_equal(result, [[0.3, 0.0, 0.7]]) # MultiObjective: one row per sub-objective. @@ -560,7 +571,7 @@ def test_extract_objective_weight_matrix(self) -> None: Objective(metric=m3, minimize=True), ] ) - result = extract_objective_weight_matrix(multi, outcomes) + result = extract_objective_weight_matrix(multi, outcomes, experiment) np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0]]) def test_get_fresh_pairwise_trial_indices(self) -> None: diff --git a/ax/adapter/tests/test_base_adapter.py b/ax/adapter/tests/test_base_adapter.py index 6e4e5eed757..bf4a72949e8 100644 --- a/ax/adapter/tests/test_base_adapter.py +++ b/ax/adapter/tests/test_base_adapter.py @@ -323,6 +323,7 @@ def test_gen_base(self, mock_fit: Mock, mock_gen_arms: Mock) -> None: ) def test_gen_on_experiment_with_imm_ss_and_opt_conf(self, _) -> None: exp = get_experiment_for_value() + exp.add_tracking_metric(Metric("test_metric")) exp._properties[Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF] = True exp.optimization_config = get_optimization_config_no_constraints() adapter = Adapter(experiment=exp, generator=Generator()) @@ -593,19 +594,22 @@ def test_set_status_quo_with_multiple_observations(self) -> None: # Case 1: Experiment has an optimization config with single map key. exp = get_branin_experiment_with_timestamp_map_metric(with_status_quo=True) # Add a second map metric, and a non-map metric. + branin_map_constraint = get_map_metric("branin_map_constraint") + branin_constraint = BraninMetric( + name="branin_constraint", + param_names=["x1", "x2"], + lower_is_better=True, + ) + exp.add_tracking_metrics([branin_map_constraint, branin_constraint]) exp.optimization_config = none_throws(exp.optimization_config).clone_with_args( outcome_constraints=[ OutcomeConstraint( - metric=get_map_metric("branin_map_constraint"), + metric=branin_map_constraint, op=ComparisonOp.LEQ, bound=5.0, ), OutcomeConstraint( - metric=BraninMetric( - name="branin_constraint", - param_names=["x1", "x2"], - lower_is_better=True, - ), + metric=branin_constraint, op=ComparisonOp.LEQ, bound=5.0, ), @@ -650,7 +654,7 @@ def test_set_status_quo_with_multiple_observations(self) -> None: ), {"branin_map", "branin_map_constraint"}, ) - opt_config_metrics = set(none_throws(exp.optimization_config).metrics) + opt_config_metrics = none_throws(exp.optimization_config).metric_names self.assertEqual(call_kwargs["metrics"], opt_config_metrics) adapter_sq = none_throws(adapter.status_quo) self.assertEqual( @@ -685,11 +689,12 @@ def test_set_status_quo_with_multiple_observations(self) -> None: ) # Case 3: Experiment doesn't have an optimization config. - metrics = none_throws(exp.optimization_config).metrics.values() + metrics = list(exp.metrics.values()) exp._optimization_config = None # Attach as tracking metric to prevent data filtering. for m in metrics: - exp.add_tracking_metric(m) + if m.name not in exp.metrics: + exp.add_tracking_metric(m) with self.assertLogs(logger=logger, level="WARN") as mock_logs: adapter = Adapter( experiment=exp, @@ -1357,15 +1362,15 @@ def test_untransform_observation_features_derived_parameter_with_digits( ), ] ) + metric = GenericNoisyFunctionMetric(name="random", f=lambda _: random()) experiment = Experiment( name="test_derived_digits", search_space=search_space, is_test=True, + tracking_metrics=[metric], optimization_config=OptimizationConfig( objective=Objective( - metric=GenericNoisyFunctionMetric( - name="random", f=lambda _: random() - ), + metric=metric, minimize=True, ) ), diff --git a/ax/adapter/tests/test_cross_validation.py b/ax/adapter/tests/test_cross_validation.py index 5001bdbf608..3626ca12b9c 100644 --- a/ax/adapter/tests/test_cross_validation.py +++ b/ax/adapter/tests/test_cross_validation.py @@ -477,6 +477,7 @@ def test_has_good_opt_config_model_fit(self) -> None: has_good_fit = has_good_opt_config_model_fit( optimization_config=optimization_config, assess_model_fit_result=assess_model_fit_result, + experiment=self.experiment, ) self.assertFalse(has_good_fit) @@ -484,14 +485,15 @@ def test_has_good_opt_config_model_fit(self) -> None: optimization_config = MultiObjectiveOptimizationConfig( objective=MultiObjective( objectives=[ - Objective(Metric("m1"), minimize=False), - Objective(Metric("m2"), minimize=False), + Objective(metric=Metric("m1"), minimize=False), + Objective(metric=Metric("m2"), minimize=False), ] ) ) has_good_fit = has_good_opt_config_model_fit( optimization_config=optimization_config, assess_model_fit_result=assess_model_fit_result, + experiment=self.experiment, ) self.assertFalse(has_good_fit) @@ -505,6 +507,7 @@ def test_has_good_opt_config_model_fit(self) -> None: has_good_fit = has_good_opt_config_model_fit( optimization_config=optimization_config, assess_model_fit_result=assess_model_fit_result, + experiment=self.experiment, ) self.assertFalse(has_good_fit) diff --git a/ax/adapter/tests/test_discrete_adapter.py b/ax/adapter/tests/test_discrete_adapter.py index f96767e2ced..523cb8af0c9 100644 --- a/ax/adapter/tests/test_discrete_adapter.py +++ b/ax/adapter/tests/test_discrete_adapter.py @@ -143,9 +143,11 @@ def test_predict(self) -> None: def test_gen(self) -> None: # Test with constraints optimization_config = OptimizationConfig( - objective=Objective(Metric("m1"), minimize=True), + objective=Objective(metric=Metric("m1"), minimize=True), outcome_constraints=[ - OutcomeConstraint(Metric("m2"), ComparisonOp.GEQ, 2, False) + OutcomeConstraint( + metric=Metric("m2"), op=ComparisonOp.GEQ, bound=2, relative=False + ) ], ) with mock.patch("ax.generators.discrete_base.DiscreteGenerator.fit"): @@ -229,9 +231,11 @@ def test_gen(self) -> None: # Test validation optimization_config = OptimizationConfig( - objective=Objective(Metric("m1"), minimize=False), + objective=Objective(metric=Metric("m1"), minimize=False), outcome_constraints=[ - OutcomeConstraint(Metric("m2"), ComparisonOp.GEQ, 2, True) + OutcomeConstraint( + metric=Metric("m2"), op=ComparisonOp.GEQ, bound=2, relative=True + ) ], ) with self.assertRaisesRegex(ValueError, "relative constraint"): diff --git a/ax/adapter/tests/test_hierarchical_search_space.py b/ax/adapter/tests/test_hierarchical_search_space.py index dd5c86ed7f6..0fd03518a3d 100644 --- a/ax/adapter/tests/test_hierarchical_search_space.py +++ b/ax/adapter/tests/test_hierarchical_search_space.py @@ -136,14 +136,14 @@ def _test_gen_base( The experiment with the generated candidates. This can be used to chain tests for other functionality that requires data. """ + metric = GenericNoisyFunctionMetric(name="random", f=lambda _: random()) experiment = Experiment( name="test_experiment", search_space=hss, + tracking_metrics=[metric], optimization_config=OptimizationConfig( objective=Objective( - metric=GenericNoisyFunctionMetric( - name="random", f=lambda _: random() - ), + metric=metric, minimize=True, ) ), diff --git a/ax/adapter/tests/test_torch_adapter.py b/ax/adapter/tests/test_torch_adapter.py index 56d52e30271..358d505b9f9 100644 --- a/ax/adapter/tests/test_torch_adapter.py +++ b/ax/adapter/tests/test_torch_adapter.py @@ -266,7 +266,9 @@ def test_TorchAdapter(self, device: torch.device | None = None) -> None: self.assertEqual(gen_opt_config.model_gen_options, {"option": "yes"}) self.assertIs(gen_opt_config.rounding_func, torch.round) self.assertFalse(gen_opt_config.is_moo) - self.assertEqual(gen_opt_config.opt_config_metrics, opt_config.metrics) + self.assertEqual( + set(gen_opt_config.opt_config_metrics.keys()), opt_config.metric_names + ) self.assertEqual(gen_args["search_space_digest"].target_values, {}) self.assertEqual(len(gen_run.arms), 1) self.assertEqual(gen_run.arms[0].parameters, {"x1": 1.0, "x2": 2.0, "x3": 3.0}) @@ -458,8 +460,10 @@ def test_best_point(self) -> None: self.assertEqual(model_predictions[1], {"m": {"m": cov}}) # test optimization config validation - raise error when - # ScalarizedOutcomeConstraint contains a metric that is not in the outcomes - with self.assertRaisesRegex(ValueError, "as a relative constraint."): + # ScalarizedOutcomeConstraint contains a metric that is not in the outcomes. + # With relative=True (default), the constraint expression contains + # "baseline" and the validation catches it as a relative constraint. + with self.assertRaisesRegex(ValueError, "relative constraint"): adapter.gen( n=1, optimization_config=OptimizationConfig( @@ -1200,7 +1204,8 @@ def test_pairwise_preference_generator(self) -> None: ), optimization_config=OptimizationConfig( Objective( - Metric(Keys.PAIRWISE_PREFERENCE_QUERY.value), minimize=False + metric=Metric(Keys.PAIRWISE_PREFERENCE_QUERY.value), + minimize=False, ) ), fit_tracking_metrics=False, diff --git a/ax/adapter/tests/test_torch_moo_adapter.py b/ax/adapter/tests/test_torch_moo_adapter.py index b01c3dbcd2e..7d01c6a893f 100644 --- a/ax/adapter/tests/test_torch_moo_adapter.py +++ b/ax/adapter/tests/test_torch_moo_adapter.py @@ -23,6 +23,7 @@ from ax.adapter.registry import Cont_X_trans, ST_MTGP_trans, Y_trans from ax.adapter.torch import TorchAdapter from ax.core.metric import Metric +from ax.core.objective import Objective from ax.core.observation import ObservationData, ObservationFeatures from ax.core.optimization_config import MultiObjectiveOptimizationConfig from ax.core.outcome_constraint import ( @@ -79,9 +80,9 @@ def helper_test_pareto_frontier( ) for trial in exp.trials.values(): trial.mark_running(no_runner_required=True).mark_completed() - metrics_dict = none_throws(exp.optimization_config).metrics + metrics_dict = exp.metrics objective_bound = 5.0 - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] = [ ObjectiveThreshold( metric=metrics_dict[f"branin_{letter}"], bound=objective_bound, @@ -90,6 +91,13 @@ def helper_test_pareto_frontier( ) for letter in "ab" ] + # Register any constraint metrics on the experiment before setting + # the optimization config. + if outcome_constraints is not None: + for oc in outcome_constraints: + for name in oc.metric_names: + if name not in exp.metrics: + exp.add_tracking_metric(Metric(name=name)) exp.optimization_config = assert_is_instance( exp.optimization_config, MultiObjectiveOptimizationConfig ).clone_with_args( @@ -243,7 +251,10 @@ def test_pareto_frontier(self) -> None: The constraint won't come close to binding, so it shouldn't affect results. """ constraint = OutcomeConstraint( - Metric(name="branin_c"), ComparisonOp.LEQ, bound=100.0, relative=False + metric=Metric(name="branin_c"), + op=ComparisonOp.LEQ, + bound=100.0, + relative=False, ) for outcome_constraints in [None, [constraint]]: with self.subTest(outcome_constraints=outcome_constraints): @@ -258,8 +269,8 @@ def test_get_pareto_frontier_and_configs_input_validation(self) -> None: ) for trial in exp.trials.values(): trial.mark_running(no_runner_required=True).mark_completed() - metrics_dict = none_throws(exp.optimization_config).metrics - objective_thresholds = [ + metrics_dict = exp.metrics + objective_thresholds: list[OutcomeConstraint] = [ ObjectiveThreshold( metric=metrics_dict[f"branin_{letter}"], bound=5.0, @@ -322,10 +333,11 @@ def test_hypervolume(self, _, cuda: bool = False) -> None: ) for trial in exp.trials.values(): trial.mark_running(no_runner_required=True).mark_completed() - metrics_dict = none_throws(exp.optimization_config).metrics + # pyre-fixme[16]: Optional type has no attribute `metrics`. + metrics_dict = exp.metrics # Objective thresholds and synthetic observations chosen to have closed-form # hypervolumes to test. - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] = [ ObjectiveThreshold( metric=metrics_dict["branin_a"], bound=10.0, @@ -354,7 +366,7 @@ def test_hypervolume(self, _, cuda: bool = False) -> None: exp.attach_data( get_branin_data_multi_objective( trial_indices=exp.trials.keys(), - outcomes=list(optimization_config.metrics), + outcomes=list(optimization_config.metric_names), ) ) adapter = TorchAdapter( @@ -447,7 +459,12 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None: param_constraints = [ParameterConstraint(inequality="x1 <= 10")] search_space.add_parameter_constraints(param_constraints) oc = none_throws(exp.optimization_config).clone() - oc.objective._objectives[0].minimize = True + # Set the first sub-objective to minimize (ensure negative sign). + sub_exprs = [s.strip() for s in oc.objective.expression.split(",")] + first = sub_exprs[0] + if not first.startswith("-"): + sub_exprs[0] = f"-{first}" + oc.objective = Objective(expression=", ".join(sub_exprs)) for use_partial_thresholds in (False, True): if use_partial_thresholds: @@ -455,7 +472,7 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None: oc, MultiObjectiveOptimizationConfig )._objective_thresholds = [ ObjectiveThreshold( - metric=oc.objective.metrics[0], + metric=Metric(name=oc.objective.metric_names[0]), bound=2.0, relative=False, op=ComparisonOp.LEQ, @@ -536,10 +553,10 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None: self.assertTrue( torch.equal(ckwargs["objective_weights"], expected_obj_weights) ) - self.assertEqual(obj_thresholds[0].metric.name, "branin_a") - self.assertEqual(obj_thresholds[1].metric.name, "branin_b") - self.assertEqual(obj_thresholds[0].metric.signature, "branin_a") - self.assertEqual(obj_thresholds[1].metric.signature, "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") self.assertEqual(obj_thresholds[0].op, ComparisonOp.LEQ) self.assertEqual(obj_thresholds[1].op, ComparisonOp.LEQ) self.assertFalse(obj_thresholds[0].relative) @@ -615,10 +632,10 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None: ckwargs = mock_get_X_pending_and_observed.call_args.kwargs self.assertEqual(ckwargs["fixed_features"], {2: 1.0}) mock_untransform_objective_thresholds.assert_called_once() - self.assertEqual(obj_thresholds[0].metric.name, "branin_a") - self.assertEqual(obj_thresholds[1].metric.name, "branin_b") - self.assertEqual(obj_thresholds[0].metric.signature, "branin_a") - self.assertEqual(obj_thresholds[1].metric.signature, "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") self.assertEqual(obj_thresholds[0].op, ComparisonOp.LEQ) self.assertEqual(obj_thresholds[1].op, ComparisonOp.LEQ) self.assertFalse(obj_thresholds[0].relative) @@ -668,10 +685,10 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None: ) mock_untransform.assert_called_once() self.assertEqual(wrapped_cast.call_count, 0) - self.assertEqual(obj_thresholds[0].metric.name, "branin_a") - self.assertEqual(obj_thresholds[1].metric.name, "branin_b") - self.assertEqual(obj_thresholds[0].metric.signature, "branin_a") - self.assertEqual(obj_thresholds[1].metric.signature, "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") self.assertEqual(obj_thresholds[0].op, ComparisonOp.LEQ) self.assertEqual(obj_thresholds[1].op, ComparisonOp.LEQ) self.assertFalse(obj_thresholds[0].relative) @@ -700,10 +717,10 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None: optimization_config=exp.optimization_config, fixed_features=None, ) - self.assertEqual(obj_thresholds[0].metric.name, "branin_a") - self.assertEqual(obj_thresholds[1].metric.name, "branin_b") - self.assertEqual(obj_thresholds[0].metric.signature, "branin_a") - self.assertEqual(obj_thresholds[1].metric.signature, "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") + self.assertEqual(obj_thresholds[0].metric_names[0], "branin_a") + self.assertEqual(obj_thresholds[1].metric_names[0], "branin_b") self.assertEqual(obj_thresholds[0].op, ComparisonOp.LEQ) self.assertEqual(obj_thresholds[1].op, ComparisonOp.LEQ) self.assertFalse(obj_thresholds[0].relative) diff --git a/ax/adapter/tests/test_transform_utils.py b/ax/adapter/tests/test_transform_utils.py index 12d65b322e9..1a9e7066245 100644 --- a/ax/adapter/tests/test_transform_utils.py +++ b/ax/adapter/tests/test_transform_utils.py @@ -33,5 +33,5 @@ def test_derelativize_optimization_config_with_raw_status_quo(self) -> None: expected_bound_values = {"m1": 0.9975, "m2": 1.995, "m3": 5.985} for oc in new_opt_config.all_constraints: self.assertFalse(oc.relative) - expected_bound_value = expected_bound_values[oc.metric.name] + expected_bound_value = expected_bound_values[oc.metric_names[0]] self.assertEqual(oc.bound, expected_bound_value) diff --git a/ax/adapter/tests/test_utils.py b/ax/adapter/tests/test_utils.py index 30853afd2d5..9664ef9579e 100644 --- a/ax/adapter/tests/test_utils.py +++ b/ax/adapter/tests/test_utils.py @@ -18,6 +18,7 @@ from ax.adapter.registry import Generators from ax.core.arm import Arm from ax.core.evaluations_to_data import raw_evaluations_to_data +from ax.core.experiment import Experiment from ax.core.generator_run import GeneratorRun from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective @@ -27,6 +28,7 @@ OutcomeConstraint, ScalarizedOutcomeConstraint, ) +from ax.core.search_space import SearchSpace from ax.core.types import ComparisonOp from ax.core.utils import get_pending_observation_features from ax.utils.common.constants import Keys @@ -78,13 +80,17 @@ def setUp(self) -> None: def test_extract_outcome_constraints(self) -> None: outcomes = ["m1", "m2", "m3"] + experiment = Experiment( + search_space=SearchSpace(parameters=[]), + tracking_metrics=[Metric(name) for name in outcomes], + ) # pass no outcome constraints - self.assertIsNone(extract_outcome_constraints([], outcomes)) + self.assertIsNone(extract_outcome_constraints([], outcomes, experiment)) outcome_constraints = [ OutcomeConstraint(metric=Metric("m1"), op=ComparisonOp.LEQ, bound=0) ] - res = extract_outcome_constraints(outcome_constraints, outcomes) + res = extract_outcome_constraints(outcome_constraints, outcomes, experiment) self.assertEqual(res[0].shape, (1, 3)) self.assertListEqual(list(res[0][0]), [1, 0, 0]) self.assertEqual(res[1][0][0], 0) @@ -98,7 +104,7 @@ def test_extract_outcome_constraints(self) -> None: bound=1, ), ] - res = extract_outcome_constraints(outcome_constraints, outcomes) + res = extract_outcome_constraints(outcome_constraints, outcomes, experiment) self.assertEqual(res[0].shape, (2, 3)) self.assertListEqual(list(res[0][0]), [1, 0, 0]) self.assertListEqual(list(res[0][1]), [0, -0.5, -0.5]) @@ -107,6 +113,10 @@ def test_extract_outcome_constraints(self) -> None: def test_extract_objective_thresholds(self) -> None: outcomes = ["m1", "m2", "m3", "m4"] + experiment = Experiment( + search_space=SearchSpace(parameters=[]), + tracking_metrics=[Metric(name) for name in outcomes], + ) objective = MultiObjective( objectives=[ Objective(metric=Metric(name), minimize=False) for name in outcomes[:3] @@ -125,7 +135,10 @@ def test_extract_objective_thresholds(self) -> None: # None of no thresholds self.assertIsNone( extract_objective_thresholds( - objective_thresholds=[], objective=objective, outcomes=outcomes + objective_thresholds=[], + objective=objective, + outcomes=outcomes, + experiment=experiment, ) ) @@ -134,6 +147,7 @@ def test_extract_objective_thresholds(self) -> None: objective_thresholds=objective_thresholds, objective=objective, outcomes=outcomes, + experiment=experiment, ) expected_obj_t_not_nan = np.array([2.0, 3.0, 4.0]) self.assertTrue(np.array_equal(obj_t[:3], expected_obj_t_not_nan[:3])) @@ -145,17 +159,19 @@ def test_extract_objective_thresholds(self) -> None: objective_thresholds=objective_thresholds[:2], objective=objective, outcomes=outcomes, + experiment=experiment, ) self.assertTrue(np.array_equal(obj_t[:2], expected_obj_t_not_nan[:2])) self.assertTrue(np.isnan(obj_t[-2:]).all()) # Fails if a threshold does not have a corresponding metric. - objective2 = Objective(Metric("m1"), minimize=False) + objective2 = Objective(expression="m1") with self.assertRaisesRegex(ValueError, "corresponding metrics"): extract_objective_thresholds( objective_thresholds=objective_thresholds, objective=objective2, outcomes=outcomes, + experiment=experiment, ) # Works with a single objective, single threshold @@ -163,6 +179,7 @@ def test_extract_objective_thresholds(self) -> None: objective_thresholds=objective_thresholds[:1], objective=objective2, outcomes=outcomes, + experiment=experiment, ) self.assertEqual(obj_t[0], 2.0) self.assertTrue(np.all(np.isnan(obj_t[1:]))) @@ -177,6 +194,7 @@ def test_extract_objective_thresholds(self) -> None: objective_thresholds=objective_thresholds, objective=objective, outcomes=outcomes, + experiment=experiment, ) objective_thresholds[2] = ObjectiveThreshold( metric=Metric("m3"), op=ComparisonOp.LEQ, bound=3, relative=True @@ -186,6 +204,7 @@ def test_extract_objective_thresholds(self) -> None: objective_thresholds=objective_thresholds, objective=objective, outcomes=outcomes, + experiment=experiment, ) def test_observation_data_to_array(self) -> None: diff --git a/ax/adapter/torch.py b/ax/adapter/torch.py index d72295e90d0..c27f646de7a 100644 --- a/ax/adapter/torch.py +++ b/ax/adapter/torch.py @@ -1018,13 +1018,18 @@ def _get_transformed_model_gen_args( "to be specified" ) - validate_transformed_optimization_config(optimization_config, self.outcomes) + validate_transformed_optimization_config( + optimization_config, self.outcomes, experiment=self._experiment + ) objective_weights = extract_objective_weight_matrix( - objective=optimization_config.objective, outcomes=self.outcomes + objective=optimization_config.objective, + outcomes=self.outcomes, + experiment=self._experiment, ) outcome_constraints = extract_outcome_constraints( outcome_constraints=optimization_config.outcome_constraints, outcomes=self.outcomes, + experiment=self._experiment, ) pruning_target_point = arm_to_np_array( arm=optimization_config.pruning_target_parameterization, @@ -1040,10 +1045,14 @@ def _get_transformed_model_gen_args( objective_thresholds=optimization_config.objective_thresholds, objective=optimization_config.objective, outcomes=self.outcomes, + experiment=self._experiment, ) else: objective_thresholds = None - opt_config_metrics = optimization_config.metrics + opt_config_metrics = { + name: self._experiment.get_metric(name) + for name in optimization_config.metric_names + } pending_array = pending_observations_as_array_list( pending_observations, self.outcomes, self.parameters @@ -1225,7 +1234,7 @@ def _validate_preference_metric_ordering( pe_aux_exp.experiment.search_space.parameters.keys() ) - pref_opt_metrics = [m.name for m in optimization_config.objective.metrics] + pref_opt_metrics = optimization_config.objective.metric_names # Get outcome order from the fitted surrogate. We must use surrogate.outcomes # which excludes auxiliary datasets (e.g., pairwise preference queries). @@ -1261,13 +1270,16 @@ def _validate_preference_metric_ordering( def validate_transformed_optimization_config( - optimization_config: OptimizationConfig, outcomes: list[str] + optimization_config: OptimizationConfig, + outcomes: list[str], + experiment: Experiment, ) -> None: """Validate optimization config against generator fitted outcomes. Args: optimization_config: Config to validate. outcomes: List of metric names w/ valid generator fits. + experiment: The experiment to look up metrics from. Raises if: 1. In the modeling layer, absolute constraints are required, however, @@ -1287,18 +1299,24 @@ def validate_transformed_optimization_config( " applied but was not." ) if isinstance(c, ScalarizedOutcomeConstraint): - for c_metric in c.metrics: + for c_name in c.metric_names: + c_metric = experiment.get_metric(c_name) if c_metric.signature not in outcomes: raise DataRequiredError( f"Scalarized constraint metric component " - f"{c.metric.signature} not found in fitted data." + f"{c_metric.signature} not found in fitted data." ) - elif c.metric.signature not in outcomes: - raise DataRequiredError( - f"Outcome constraint metric {c.metric.signature} not found in fitted " - "data." - ) - obj_metric_signatures = [m.signature for m in optimization_config.objective.metrics] + else: + c_metric = experiment.get_metric(c.metric_names[0]) + if c_metric.signature not in outcomes: + raise DataRequiredError( + f"Outcome constraint metric {c_metric.signature} not found in " + "fitted data." + ) + obj_metric_signatures = [ + experiment.get_metric(name).signature + for name in optimization_config.objective.metric_names + ] for obj_metric_signature in obj_metric_signatures: if obj_metric_signature not in outcomes: raise DataRequiredError( diff --git a/ax/adapter/transforms/base.py b/ax/adapter/transforms/base.py index 87ba7472fa7..009b7c8d5a4 100644 --- a/ax/adapter/transforms/base.py +++ b/ax/adapter/transforms/base.py @@ -87,6 +87,22 @@ def __init__( self.config = deepcopy(config) self.adapter = adapter + def _get_metric_signature( + self, + metric_name: str, + adapter: adapter_module.base.Adapter | None = None, + ) -> str: + """Get the metric signature for a metric name. + + Uses the experiment from the adapter to look up the actual signature. + Falls back to using the metric name directly if no adapter is available + (which is correct for simple Metric objects where name == signature). + """ + a = adapter or self.adapter + if a is not None: + return a._experiment.get_metric(metric_name).signature + return metric_name + def transform_search_space(self, search_space: SearchSpace) -> SearchSpace: """Transform search space. diff --git a/ax/adapter/transforms/bilog_y.py b/ax/adapter/transforms/bilog_y.py index ee68a53f004..7ef2f240b57 100644 --- a/ax/adapter/transforms/bilog_y.py +++ b/ax/adapter/transforms/bilog_y.py @@ -68,9 +68,10 @@ def __init__( if adapter is not None and adapter._optimization_config is not None: # TODO @deriksson: Add support for relative outcome constraints self.metric_to_bound: dict[str, float] = { - oc.metric.signature: oc.bound + metric.signature: oc.bound for oc in adapter._optimization_config.outcome_constraints if not oc.relative + for metric in adapter._experiment.get_metrics(oc.metric_names) } else: self.metric_to_bound = {} diff --git a/ax/adapter/transforms/derelativize.py b/ax/adapter/transforms/derelativize.py index 20c638a0489..3f94861c052 100644 --- a/ax/adapter/transforms/derelativize.py +++ b/ax/adapter/transforms/derelativize.py @@ -15,12 +15,17 @@ from ax.adapter.observation_utils import unwrap_observation_data from ax.adapter.transforms.base import Transform from ax.adapter.transforms.ivw import ivw_metric_merge +from ax.core.experiment import Experiment from ax.core.observation import ObservationFeatures -from ax.core.optimization_config import OptimizationConfig +from ax.core.optimization_config import ( + MultiObjectiveOptimizationConfig, + OptimizationConfig, +) from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint -from ax.core.types import TModelCov, TModelMean +from ax.core.types import ComparisonOp, TModelCov, TModelMean from ax.exceptions.core import DataRequiredError from ax.utils.common.logger import get_logger +from ax.utils.common.sympy import build_constraint_expression_str from pyre_extensions import none_throws @@ -89,41 +94,46 @@ def transform_optimization_config( raw_f=raw_f, pred_f=f, pred_cov=cov, + experiment=adapter._experiment, ) else: f = raw_f # Plug in the status quo value to each relative constraint. - for c in optimization_config.all_constraints: + # Build new constraint lists since constraints are immutable. + new_outcome_constraints = [] + for c in optimization_config.outcome_constraints: if c.relative: - if isinstance(c, ScalarizedOutcomeConstraint): - missing_metrics = { - metric.signature - for metric in c.metrics - if metric.signature not in f - } - if len(missing_metrics) > 0: - raise DataRequiredError( - f"Status-quo metric value not yet available for metric(s) " - f"{missing_metrics}." - ) - # The sq_val of scalarized outcome is the weighted - # sum of its component metrics - sq_val = np.sum( - [ - c.weights[i] * f[metric.signature][0] - for i, metric in enumerate(c.metrics) - ] + sq_val = _compute_sq_val(c, adapter, f) + new_bound = derelativize_bound(bound=c.bound, sq_val=sq_val) + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=new_bound, + relative=False, + ) + new_outcome_constraints.append(OutcomeConstraint(expression=new_expr)) + else: + new_outcome_constraints.append(c) + optimization_config.outcome_constraints = new_outcome_constraints + + if isinstance(optimization_config, MultiObjectiveOptimizationConfig): + new_thresholds = [] + for c in optimization_config.objective_thresholds: + if c.relative: + sq_val = _compute_sq_val(c, adapter, f) + new_bound = derelativize_bound(bound=c.bound, sq_val=sq_val) + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=new_bound, + relative=False, ) - elif c.metric.signature in f: - sq_val = f[c.metric.signature][0] + new_thresholds.append(OutcomeConstraint(expression=new_expr)) else: - raise DataRequiredError( - f"Status-quo metric value not yet available for metric " - f"{c.metric.name}." - ) - c.bound = derelativize_bound(bound=c.bound, sq_val=sq_val) - c.relative = False + new_thresholds.append(c) + optimization_config.objective_thresholds = new_thresholds + return optimization_config def untransform_outcome_constraints( @@ -136,25 +146,68 @@ def untransform_outcome_constraints( return outcome_constraints +def _compute_sq_val( + c: OutcomeConstraint, + adapter: adapter_module.base.Adapter, + f: TModelMean, +) -> float: + """Compute the status quo value for a constraint. + + For scalarized constraints, returns the weighted sum of component metric + status quo values. For single-metric constraints, returns the status quo + value for that metric. + """ + if isinstance(c, ScalarizedOutcomeConstraint): + metrics = adapter._experiment.get_metrics(c.metric_names) + missing_metrics = { + metric.signature for metric in metrics if metric.signature not in f + } + if len(missing_metrics) > 0: + raise DataRequiredError( + f"Status-quo metric value not yet available for metric(s) " + f"{missing_metrics}." + ) + # The sq_val of scalarized outcome is the weighted + # sum of its component metrics + return float( + np.sum( + [ + c.weights[i] * f[metric.signature][0] + for i, metric in enumerate(metrics) + ] + ) + ) + else: + metric = adapter._experiment.get_metrics(c.metric_names)[0] + if metric.signature in f: + return f[metric.signature][0] + else: + raise DataRequiredError( + f"Status-quo metric value not yet available for metric " + f"{c.metric_names[0]}." + ) + + def _warn_if_raw_sq_is_out_of_CI( optimization_config: OptimizationConfig, raw_f: TModelMean, pred_f: TModelMean, pred_cov: TModelCov, + experiment: Experiment, ) -> None: """Warn if the raw SQ values for relative constraint metrics deviate by more than 1.96 standard deviation from the predictions. """ relative_metrics = { - oc.metric.signature + experiment.get_metric(oc.metric_names[0]).signature for oc in optimization_config.all_constraints if oc.relative and not isinstance(oc, ScalarizedOutcomeConstraint) }.union( { - metric.signature + experiment.get_metric(name).signature for oc in optimization_config.all_constraints if oc.relative and isinstance(oc, ScalarizedOutcomeConstraint) - for metric in oc.metrics + for name in oc.metric_names } ) for metric_signature in relative_metrics: diff --git a/ax/adapter/transforms/log_y.py b/ax/adapter/transforms/log_y.py index 4eec60ae05b..03f59219c9d 100644 --- a/ax/adapter/transforms/log_y.py +++ b/ax/adapter/transforms/log_y.py @@ -18,11 +18,16 @@ from ax.adapter.transforms.base import Transform from ax.adapter.transforms.utils import match_ci_width, T_MATCH_CI_WIDTH from ax.core.observation import ObservationData, ObservationFeatures -from ax.core.optimization_config import OptimizationConfig +from ax.core.optimization_config import ( + MultiObjectiveOptimizationConfig, + OptimizationConfig, +) from ax.core.outcome_constraint import OutcomeConstraint from ax.core.search_space import SearchSpace +from ax.core.types import ComparisonOp from ax.generators.types import TConfig from ax.utils.common.logger import get_logger +from ax.utils.common.sympy import build_constraint_expression_str from pyre_extensions import assert_is_instance @@ -86,9 +91,13 @@ def transform_optimization_config( adapter: base_adapter.Adapter | None = None, fixed_features: ObservationFeatures | None = None, ) -> OptimizationConfig: - for c in optimization_config.all_constraints: - if c.metric.signature in self.metric_signatures: - base_str = f"LogY transform cannot be applied to metric {c.metric.name}" + new_outcome_constraints = [] + for c in optimization_config.outcome_constraints: + sig = self._get_metric_signature(c.metric_names[0], adapter) + if sig in self.metric_signatures: + base_str = ( + f"LogY transform cannot be applied to metric {c.metric_names[0]}" + ) if c.relative: raise ValueError( f"{base_str} since it is subject to a relative constraint." @@ -98,7 +107,49 @@ def transform_optimization_config( f"{base_str} since the bound isn't positive, got: {c.bound}." ) else: - c.bound = np.log(c.bound) + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=np.log(c.bound), + relative=c.relative, + ) + new_outcome_constraints.append( + OutcomeConstraint(expression=new_expr) + ) + else: + new_outcome_constraints.append(c) + optimization_config.outcome_constraints = new_outcome_constraints + + if isinstance(optimization_config, MultiObjectiveOptimizationConfig): + new_thresholds = [] + for c in optimization_config.objective_thresholds: + sig = self._get_metric_signature(c.metric_names[0], adapter) + if sig in self.metric_signatures: + base_str = ( + f"LogY transform cannot be applied to metric " + f"{c.metric_names[0]}" + ) + if c.relative: + raise ValueError( + f"{base_str} since it is subject to a relative constraint." + ) + elif c.bound <= 0: + raise ValueError( + f"{base_str} since the bound isn't positive, " + f"got: {c.bound}." + ) + else: + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=np.log(c.bound), + relative=c.relative, + ) + new_thresholds.append(OutcomeConstraint(expression=new_expr)) + else: + new_thresholds.append(c) + optimization_config.objective_thresholds = new_thresholds + return optimization_config def _tf_obs_data( @@ -156,12 +207,22 @@ def untransform_outcome_constraints( outcome_constraints: list[OutcomeConstraint], fixed_features: ObservationFeatures | None = None, ) -> list[OutcomeConstraint]: + result = [] for c in outcome_constraints: - if c.metric.signature in self.metric_signatures: + sig = self._get_metric_signature(c.metric_names[0]) + if sig in self.metric_signatures: if c.relative: raise ValueError("Unexpected relative transform.") - c.bound = np.exp(c.bound) - return outcome_constraints + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=np.exp(c.bound), + relative=c.relative, + ) + result.append(OutcomeConstraint(expression=new_expr)) + else: + result.append(c) + return result def transform_experiment_data( self, experiment_data: ExperimentData diff --git a/ax/adapter/transforms/objective_as_constraint.py b/ax/adapter/transforms/objective_as_constraint.py index e7d16e179dc..115512eed97 100644 --- a/ax/adapter/transforms/objective_as_constraint.py +++ b/ax/adapter/transforms/objective_as_constraint.py @@ -14,7 +14,7 @@ import pandas as pd from ax.adapter.data_utils import ExperimentData from ax.adapter.transforms.base import Transform -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import ScalarizedObjective from ax.core.observation import ObservationData, ObservationFeatures from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, @@ -31,6 +31,7 @@ if TYPE_CHECKING: # import as module to make sphinx-autodoc-typehints happy from ax import adapter as adapter_module # noqa F401 + from ax.core.experiment import Experiment logger: logging.Logger = get_logger(__name__) @@ -156,6 +157,7 @@ def _check_no_feasible_points( row=row, constraints=outcome_constraints, sq_data=sq_data, + experiment=adapter._experiment, ): return False @@ -194,15 +196,21 @@ def transform_optimization_config( # Handle ScalarizedObjective: create a single ScalarizedOutcomeConstraint # with the bound equal to the status quo value of the scalarized objective. if isinstance(objective, ScalarizedObjective): + obj_metrics = none_throws(adapter or self.adapter)._experiment.get_metrics( + objective.metric_names + ) scalarized_sq_value = 0.0 - for metric, weight in objective.metric_weights: + for metric, (_, weight) in zip(obj_metrics, objective.metric_weights): metric_idx = sq_data.metric_signatures.index(metric.signature) scalarized_sq_value += weight * sq_data.means[metric_idx] - op = ComparisonOp.LEQ if objective.minimize else ComparisonOp.GEQ + # metric_weights are sign-encoded (negative = minimize direction). + # With sign-encoded weights, "don't get worse than SQ" is always + # GEQ: Σ(signed_wi * yi) >= Σ(signed_wi * sq_yi). + op = ComparisonOp.GEQ new_constraint = ScalarizedOutcomeConstraint( - metrics=[m.clone() for m in objective.metrics], - weights=list(objective.weights), + metrics=[m.clone() for m in obj_metrics], + weights=[w for _, w in objective.metric_weights], op=op, bound=float(scalarized_sq_value), relative=False, @@ -211,26 +219,22 @@ def transform_optimization_config( self._scalarized_objective_constraint_added = True return optimization_config - # Get list of objectives to add constraints for. - if isinstance(optimization_config, MultiObjectiveOptimizationConfig): - # MultiObjectiveOptimizationConfig can have MultiObjective or - # ScalarizedObjective. Only MultiObjective has multiple objectives. - if isinstance(objective, MultiObjective): - objectives = objective.objectives - else: - objectives = [ - Objective(metric=objective.metric, minimize=objective.minimize) - ] - else: - objectives = [objective] + # Get list of (metric_name, weight) pairs for each objective metric. + # For multi-objective, each sub-expression has one metric. + # For single-objective, metric_weights has one entry. + obj_metric_weights = objective.metric_weights # Add a constraint for each objective at the status quo value. - for obj in objectives: - metric = obj.metric + for metric_name, weight in obj_metric_weights: + metric = none_throws(adapter or self.adapter)._experiment.get_metrics( + [metric_name] + )[0] metric_idx = sq_data.metric_signatures.index(metric.signature) sq_value = sq_data.means[metric_idx] - op = ComparisonOp.LEQ if obj.minimize else ComparisonOp.GEQ + # Negative weight in the maximize-expression means minimize + is_minimize = weight < 0 + op = ComparisonOp.LEQ if is_minimize else ComparisonOp.GEQ new_constraint = OutcomeConstraint( metric=metric.clone(), op=op, @@ -239,7 +243,7 @@ def transform_optimization_config( ) optimization_config._outcome_constraints.append(new_constraint) - self._objective_metrics_added.append(metric.name) + self._objective_metrics_added.append(metric_name) return optimization_config @@ -258,7 +262,7 @@ def untransform_outcome_constraints( outcome_constraints = [ oc for oc in outcome_constraints - if oc.metric.name not in self._objective_metrics_added + if oc.metric_names[0] not in self._objective_metrics_added ] return outcome_constraints @@ -267,6 +271,7 @@ def _is_point_feasible( row: pd.Series, constraints: list[OutcomeConstraint], sq_data: ObservationData | None = None, + experiment: Experiment | None = None, ) -> bool: """Check if a single observation satisfies all outcome constraints. @@ -280,12 +285,18 @@ def _is_point_feasible( sq_data: Status quo observation data, required for evaluating relative constraints. If None and a relative constraint is encountered, the constraint is skipped. + experiment: The experiment, used to look up metric signatures from + metric names. If None, the metric name is used as the signature. Returns: True if the point satisfies all constraints, False otherwise. """ for constraint in constraints: - metric_sig = constraint.metric.signature + metric_name = constraint.metric_names[0] + if experiment is not None: + metric_sig = experiment.get_metric(metric_name).signature + else: + metric_sig = metric_name try: mean_val = row["mean", metric_sig] except KeyError: diff --git a/ax/adapter/transforms/power_transform_y.py b/ax/adapter/transforms/power_transform_y.py index c44d7e6220f..37b75c72046 100644 --- a/ax/adapter/transforms/power_transform_y.py +++ b/ax/adapter/transforms/power_transform_y.py @@ -15,12 +15,16 @@ from ax.adapter.data_utils import ExperimentData from ax.adapter.transforms.base import Transform from ax.adapter.transforms.utils import match_ci_width -from ax.core.objective import ScalarizedObjective from ax.core.observation import ObservationData, ObservationFeatures -from ax.core.optimization_config import OptimizationConfig -from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint +from ax.core.optimization_config import ( + MultiObjectiveOptimizationConfig, + OptimizationConfig, +) +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.search_space import SearchSpace +from ax.core.types import ComparisonOp from ax.generators.types import TConfig +from ax.utils.common.sympy import build_constraint_expression_str from ax.utils.common.typeutils import assert_is_instance_list from pyre_extensions import assert_is_instance, none_throws from sklearn.preprocessing import PowerTransformer @@ -152,13 +156,12 @@ def transform_optimization_config( adapter: adapter_module.base.Adapter | None = None, fixed_features: ObservationFeatures | None = None, ) -> OptimizationConfig: - if isinstance(optimization_config.objective, ScalarizedObjective): - objective_metric_signatures = [ - metric.signature for metric in optimization_config.objective.metrics + if optimization_config.objective.is_scalarized_objective: + objective_metric_sigs = [ + self._get_metric_signature(n, adapter) + for n in optimization_config.objective.metric_names ] - intersection = set(objective_metric_signatures) & set( - self.metric_signatures - ) + intersection = set(objective_metric_sigs) & set(self.metric_signatures) if intersection: raise NotImplementedError( "PowerTransformY cannot be used for metric(s) " @@ -167,26 +170,85 @@ def transform_optimization_config( "preserve the linear scalarization of the objective." ) - for c in optimization_config.all_constraints: - if isinstance(c, ScalarizedOutcomeConstraint): - c_metric_signatures = [metric.signature for metric in c.metrics] - intersection = set(c_metric_signatures) & set(self.metric_signatures) + new_outcome_constraints = [] + for c in optimization_config.outcome_constraints: + if len(c.metric_names) > 1: + c_sigs = [ + self._get_metric_signature(n, adapter) for n in c.metric_names + ] + intersection = set(c_sigs) & set(self.metric_signatures) if intersection: raise NotImplementedError( "PowerTransformY cannot be used for metric(s) " f"{intersection} that are part of a " "ScalarizedOutcomeConstraint." ) - elif c.metric.signature in self.metric_signatures: - if c.relative: - raise ValueError( - "PowerTransformY cannot be applied to metric " - f"{c.metric.signature} since it is subject to " - "a relative constraint." - ) + new_outcome_constraints.append(c) + else: + c_sig = self._get_metric_signature(c.metric_names[0], adapter) + if c_sig in self.metric_signatures: + if c.relative: + raise ValueError( + "PowerTransformY cannot be applied to metric " + f"{c_sig} since it is subject to " + "a relative constraint." + ) + else: + transform = self.power_transforms[c_sig].transform + new_bound = transform(np.array(c.bound, ndmin=2)).item() + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=new_bound, + relative=c.relative, + ) + new_outcome_constraints.append( + OutcomeConstraint(expression=new_expr) + ) else: - transform = self.power_transforms[c.metric.signature].transform - c.bound = transform(np.array(c.bound, ndmin=2)).item() + new_outcome_constraints.append(c) + optimization_config.outcome_constraints = new_outcome_constraints + + if isinstance(optimization_config, MultiObjectiveOptimizationConfig): + new_thresholds = [] + for c in optimization_config.objective_thresholds: + if len(c.metric_names) > 1: + c_sigs = [ + self._get_metric_signature(n, adapter) for n in c.metric_names + ] + intersection = set(c_sigs) & set(self.metric_signatures) + if intersection: + raise NotImplementedError( + "PowerTransformY cannot be used for metric(s) " + f"{intersection} that are part of a " + "ScalarizedOutcomeConstraint." + ) + new_thresholds.append(c) + else: + c_sig = self._get_metric_signature(c.metric_names[0], adapter) + if c_sig in self.metric_signatures: + if c.relative: + raise ValueError( + "PowerTransformY cannot be applied to metric " + f"{c_sig} since it is subject to " + "a relative constraint." + ) + else: + transform = self.power_transforms[c_sig].transform + new_bound = transform(np.array(c.bound, ndmin=2)).item() + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=new_bound, + relative=c.relative, + ) + new_thresholds.append( + OutcomeConstraint(expression=new_expr) + ) + else: + new_thresholds.append(c) + optimization_config.objective_thresholds = new_thresholds + return optimization_config def untransform_outcome_constraints( @@ -194,18 +256,28 @@ def untransform_outcome_constraints( outcome_constraints: list[OutcomeConstraint], fixed_features: ObservationFeatures | None = None, ) -> list[OutcomeConstraint]: + result = [] for c in outcome_constraints: - if isinstance(c, ScalarizedOutcomeConstraint): + if len(c.metric_names) > 1: raise ValueError("ScalarizedOutcomeConstraint not supported here") - elif c.metric.signature in self.metric_signatures: - if c.relative: - raise ValueError("Relative constraints not supported here.") + else: + c_sig = self._get_metric_signature(c.metric_names[0]) + if c_sig in self.metric_signatures: + if c.relative: + raise ValueError("Relative constraints not supported here.") + else: + transform = self.power_transforms[c_sig].inverse_transform + new_bound = transform(np.array(c.bound, ndmin=2)).item() + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=new_bound, + relative=c.relative, + ) + result.append(OutcomeConstraint(expression=new_expr)) else: - transform = self.power_transforms[ - c.metric.signature - ].inverse_transform - c.bound = transform(np.array(c.bound, ndmin=2)).item() - return outcome_constraints + result.append(c) + return result def transform_experiment_data( self, experiment_data: ExperimentData diff --git a/ax/adapter/transforms/relativize.py b/ax/adapter/transforms/relativize.py index d7d4c041f42..4ecdb518c32 100644 --- a/ax/adapter/transforms/relativize.py +++ b/ax/adapter/transforms/relativize.py @@ -18,7 +18,6 @@ import numpy.typing as npt from ax.adapter.data_utils import ExperimentData from ax.adapter.transforms.base import Transform -from ax.core.objective import MultiObjective from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, @@ -27,8 +26,10 @@ ) from ax.core.outcome_constraint import OutcomeConstraint from ax.core.search_space import SearchSpace +from ax.core.types import ComparisonOp from ax.exceptions.core import DataRequiredError, UnsupportedError from ax.generators.types import TConfig +from ax.utils.common.sympy import build_constraint_expression_str from ax.utils.stats.math_utils import relativize, unrelativize from pyre_extensions import none_throws @@ -37,6 +38,21 @@ from ax import adapter as adapter_module # noqa F401 +def _constraint_with_relative( + c: OutcomeConstraint, relative: bool +) -> OutcomeConstraint: + """Create a new OutcomeConstraint identical to ``c`` but with a different + ``relative`` flag. Since OutcomeConstraint is immutable, we rebuild the + expression string with the desired relativity.""" + new_expr = build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=c.bound, + relative=relative, + ) + return OutcomeConstraint(expression=new_expr) + + class BaseRelativize(Transform, ABC): """ Change the relative flag of the given relative optimization configuration @@ -107,9 +123,6 @@ def transform_optimization_config( """ # Getting constraints - constraints = [ - constraint.clone() for constraint in optimization_config.outcome_constraints - ] if not all( constraint.relative for constraint in optimization_config.outcome_constraints @@ -117,13 +130,15 @@ def transform_optimization_config( raise ValueError( "All constraints must be relative to use the Relativize transform." ) - for constraint in constraints: - constraint.relative = False + constraints = [ + _constraint_with_relative(constraint, relative=False) + for constraint in optimization_config.outcome_constraints + ] if isinstance(optimization_config, PreferenceOptimizationConfig): objective = optimization_config.objective - assert isinstance(objective, MultiObjective), ( - f"Expected MultiObjective, got {type(objective).__name__}" + assert objective.is_multi_objective, ( + "Expected multi-objective, got single-objective" ) new_optimization_config = optimization_config.clone_with_args( objective=objective, @@ -131,17 +146,16 @@ def transform_optimization_config( ) elif isinstance(optimization_config, MultiObjectiveOptimizationConfig): # Getting objective thresholds - obj_thresholds = [ - obj_threshold.clone() - for obj_threshold in optimization_config.objective_thresholds - ] - for obj_threshold in obj_thresholds: + obj_thresholds = [] + for obj_threshold in optimization_config.objective_thresholds: if not obj_threshold.relative: raise ValueError( "All objective thresholds must be relative to use " "the Relativize transform." ) - obj_threshold.relative = False + obj_thresholds.append( + _constraint_with_relative(obj_threshold, relative=False) + ) new_optimization_config = optimization_config.clone_with_args( objective=optimization_config.objective, @@ -160,9 +174,9 @@ def untransform_outcome_constraints( outcome_constraints: list[OutcomeConstraint], fixed_features: ObservationFeatures | None = None, ) -> list[OutcomeConstraint]: - for c in outcome_constraints: - c.relative = True - return outcome_constraints + return [ + _constraint_with_relative(c, relative=True) for c in outcome_constraints + ] def transform_observations( self, diff --git a/ax/adapter/transforms/standardize_y.py b/ax/adapter/transforms/standardize_y.py index 6f7a7b32b24..e5d2df2ba3e 100644 --- a/ax/adapter/transforms/standardize_y.py +++ b/ax/adapter/transforms/standardize_y.py @@ -15,16 +15,19 @@ import numpy as np from ax.adapter.data_utils import ExperimentData from ax.adapter.transforms.base import Transform -from ax.core.metric import Metric -from ax.core.objective import ScalarizedObjective +from ax.core.objective import Objective from ax.core.observation import ObservationData, ObservationFeatures -from ax.core.optimization_config import OptimizationConfig -from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint +from ax.core.optimization_config import ( + MultiObjectiveOptimizationConfig, + OptimizationConfig, +) +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.search_space import SearchSpace -from ax.core.types import TParamValue +from ax.core.types import ComparisonOp, TParamValue from ax.exceptions.core import DataRequiredError from ax.generators.types import TConfig from ax.utils.common.logger import get_logger +from ax.utils.common.sympy import build_constraint_expression_str from pyre_extensions import none_throws @@ -79,10 +82,12 @@ def _transform_observation_data( obsd.covariance /= np.dot(stds[:, None], stds[:, None].transpose()) return observation_data - def _check_metrics_available(self, metrics: list[Metric], context: str) -> set[str]: + def _check_metrics_available( + self, metric_signatures: list[str], context: str + ) -> set[str]: """Check that all metrics are available and return the set of metrics.""" available_metrics = set(self.Ymean).intersection(set(self.Ystd)) - required_metrics = {metric.signature for metric in metrics} + required_metrics = set(metric_signatures) if len(required_metrics - available_metrics) > 0: raise DataRequiredError( f"`StandardizeY` transform requires {context} metric(s) " @@ -91,7 +96,7 @@ def _check_metrics_available(self, metrics: list[Metric], context: str) -> set[s return available_metrics def _transform_scalarized_weights( - self, metrics: list[Metric], weights: list[float] + self, metric_signatures: list[str], weights: list[float] ) -> list[float]: """Transform weights for scalarized objectives/constraints. @@ -100,15 +105,15 @@ def _transform_scalarized_weights( This method returns the new weights: new_wi = wi * σi. Args: - metrics: List of metrics in the scalarized term. + metric_signatures: List of metric signatures in the scalarized term. weights: Original weights for each metric. Returns: Transformed weights scaled by standard deviations. """ return [ - weights[i] * float(self.Ystd[metric.signature]) - for i, metric in enumerate(metrics) + weights[i] * float(self.Ystd[sig]) + for i, sig in enumerate(metric_signatures) ] def transform_optimization_config( @@ -117,39 +122,97 @@ def transform_optimization_config( adapter: base_adapter.Adapter | None = None, fixed_features: ObservationFeatures | None = None, ) -> OptimizationConfig: - # Handle ScalarizedObjective - if isinstance(optimization_config.objective, ScalarizedObjective): + # Handle scalarized objective (linear combination of metrics). + if optimization_config.objective.is_scalarized_objective: objective = optimization_config.objective - self._check_metrics_available(objective.metrics, context="objective") - objective.weights = self._transform_scalarized_weights( - objective.metrics, objective.weights + obj_sigs = [ + self._get_metric_signature(n, adapter) for n in objective.metric_names + ] + self._check_metrics_available(obj_sigs, context="objective") + old_weights = [w for _, w in objective.metric_weights] + new_weights = self._transform_scalarized_weights(obj_sigs, old_weights) + new_metric_weights = [ + (name, new_w) + for (name, _), new_w in zip(objective.metric_weights, new_weights) + ] + optimization_config.objective = _build_objective_from_metric_weights( + new_metric_weights + ) + + new_constraints = self._transform_constraints( + optimization_config.outcome_constraints, adapter + ) + optimization_config.outcome_constraints = new_constraints + + # For MOO configs, transform objective thresholds separately. + if isinstance(optimization_config, MultiObjectiveOptimizationConfig): + new_thresholds = self._transform_constraints( + optimization_config.objective_thresholds, adapter ) + optimization_config.objective_thresholds = new_thresholds + + return optimization_config - for c in optimization_config.all_constraints: + def _transform_constraints( + self, + constraints: list[OutcomeConstraint], + adapter: base_adapter.Adapter | None = None, + ) -> list[OutcomeConstraint]: + """Transform a list of constraints by standardizing bounds.""" + new_constraints = [] + for c in constraints: if c.relative: raise ValueError( f"StandardizeY transform does not support relative constraint {c}" ) - if isinstance(c, ScalarizedOutcomeConstraint): - self._check_metrics_available(c.metrics, context="constraint") + if len(c.metric_names) > 1: + c_sigs = [ + self._get_metric_signature(n, adapter) for n in c.metric_names + ] + self._check_metrics_available(c_sigs, context="constraint") # Transform Σ(wi * yi) <= C to Σ(wi * σi * zi) <= C - Σ(wi * μi) # where zi = (yi - μi) / σi + old_weights = [w for _, w in c.metric_weights] agg_mean = np.sum( [ - c.weights[i] * float(self.Ymean[metric.signature]) - for i, metric in enumerate(c.metrics) + old_weights[i] * float(self.Ymean[sig]) + for i, sig in enumerate(c_sigs) ] ) - c.bound = float(c.bound - agg_mean) - c.weights = self._transform_scalarized_weights(c.metrics, c.weights) + new_bound = float(c.bound - agg_mean) + new_weights = self._transform_scalarized_weights(c_sigs, old_weights) + new_metric_weights = [ + (name, new_w) + for (name, _), new_w in zip(c.metric_weights, new_weights) + ] + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=new_metric_weights, + op=op_str, + bound=new_bound, + relative=False, + ) + ) + ) else: - self._check_metrics_available([c.metric], context="constraint") - c.bound = float( - (c.bound - self.Ymean[c.metric.signature]) - / self.Ystd[c.metric.signature] + c_sig = self._get_metric_signature(c.metric_names[0], adapter) + self._check_metrics_available([c_sig], context="constraint") + new_bound = float((c.bound - self.Ymean[c_sig]) / self.Ystd[c_sig]) + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=op_str, + bound=new_bound, + relative=c.relative, + ) + ) ) - return optimization_config + return new_constraints def _untransform_observation_data( self, @@ -167,17 +230,28 @@ def untransform_outcome_constraints( outcome_constraints: list[OutcomeConstraint], fixed_features: ObservationFeatures | None = None, ) -> list[OutcomeConstraint]: + new_constraints = [] for c in outcome_constraints: if c.relative: raise ValueError( f"StandardizeY transform does not support relative constraint {c}" ) - if isinstance(c, ScalarizedOutcomeConstraint): + if len(c.metric_names) > 1: raise ValueError("ScalarizedOutcomeConstraint not supported") - c.bound = float( - c.bound * self.Ystd[c.metric.signature] + self.Ymean[c.metric.signature] + c_sig = self._get_metric_signature(c.metric_names[0]) + new_bound = float(c.bound * self.Ystd[c_sig] + self.Ymean[c_sig]) + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=op_str, + bound=new_bound, + relative=c.relative, + ) + ) ) - return outcome_constraints + return new_constraints def transform_experiment_data( self, experiment_data: ExperimentData @@ -199,6 +273,29 @@ def transform_experiment_data( ) +def _build_objective_from_metric_weights( + metric_weights: list[tuple[str, float]], +) -> Objective: + """Build a new Objective from (metric_name, weight) pairs. + + Args: + metric_weights: List of (metric_name, weight) tuples. + + Returns: + A new Objective with the corresponding expression. + """ + parts: list[str] = [] + for name, w in metric_weights: + if w == 1.0: + parts.append(name) + elif w == -1.0: + parts.append(f"-{name}") + else: + parts.append(f"{w}*{name}") + expr_str = " + ".join(parts).replace(" + -", " - ") + return Objective(expression=expr_str) + + _TYKey = TypeVar("_TYKey", bound=str | tuple[str, TParamValue]) diff --git a/ax/adapter/transforms/stratified_standardize_y.py b/ax/adapter/transforms/stratified_standardize_y.py index c38b89e523a..fbbc54aac1f 100644 --- a/ax/adapter/transforms/stratified_standardize_y.py +++ b/ax/adapter/transforms/stratified_standardize_y.py @@ -13,15 +13,21 @@ import numpy as np from ax.adapter.data_utils import ExperimentData from ax.adapter.transforms.base import Transform -from ax.adapter.transforms.standardize_y import compute_standardization_parameters -from ax.core.objective import ScalarizedObjective +from ax.adapter.transforms.standardize_y import ( + _build_objective_from_metric_weights, + compute_standardization_parameters, +) from ax.core.observation import Observation, ObservationFeatures -from ax.core.optimization_config import OptimizationConfig -from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint +from ax.core.optimization_config import ( + MultiObjectiveOptimizationConfig, + OptimizationConfig, +) +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.parameter import ChoiceParameter from ax.core.search_space import SearchSpace -from ax.core.types import TParamValue +from ax.core.types import ComparisonOp, TParamValue from ax.generators.types import TConfig +from ax.utils.common.sympy import build_constraint_expression_str from pyre_extensions import assert_is_instance, none_throws @@ -163,8 +169,8 @@ def transform_optimization_config( adapter: adapter_module.base.Adapter | None = None, fixed_features: ObservationFeatures | None = None, ) -> OptimizationConfig: - if len(optimization_config.all_constraints) == 0 and not isinstance( - optimization_config.objective, ScalarizedObjective + if len(optimization_config.all_constraints) == 0 and not ( + optimization_config.objective.is_scalarized_objective ): return optimization_config if fixed_features is None or self.p_name not in fixed_features.parameters: @@ -175,41 +181,103 @@ def transform_optimization_config( v = none_throws(fixed_features.parameters[self.p_name]) strata = self.strata_mapping[v] - # Handle ScalarizedObjective: update weights by multiplying with std. + # Handle scalarized objective: update weights by multiplying with std. # Transform \sum (wi * yi) to \sum (wi * si * zi) where zi = (yi - mu_i) / si # The constant term \sum (wi * mu_i) doesn't affect optimization. - if isinstance(objective := optimization_config.objective, ScalarizedObjective): - objective.weights = [ - objective.weights[i] * float(self.Ystd[(metric.signature, strata)]) - for i, metric in enumerate(objective.metrics) + if optimization_config.objective.is_scalarized_objective: + objective = optimization_config.objective + obj_sigs = [ + self._get_metric_signature(n, adapter) for n in objective.metric_names ] + old_weights = [w for _, w in objective.metric_weights] + new_weights = [ + old_weights[i] * float(self.Ystd[(sig, strata)]) + for i, sig in enumerate(obj_sigs) + ] + new_metric_weights = [ + (name, new_w) + for (name, _), new_w in zip(objective.metric_weights, new_weights) + ] + optimization_config.objective = _build_objective_from_metric_weights( + new_metric_weights + ) + + optimization_config.outcome_constraints = self._transform_constraints( + optimization_config.outcome_constraints, strata, adapter + ) + + if isinstance(optimization_config, MultiObjectiveOptimizationConfig): + optimization_config.objective_thresholds = self._transform_constraints( + optimization_config.objective_thresholds, strata, adapter + ) - for c in optimization_config.all_constraints: + return optimization_config + + def _transform_constraints( + self, + constraints: list[OutcomeConstraint], + strata: TParamValue, + adapter: adapter_module.base.Adapter | None = None, + ) -> list[OutcomeConstraint]: + """Transform a list of constraints by standardizing bounds.""" + new_constraints = [] + for c in constraints: if c.relative: raise ValueError( "StratifiedStandardizeY transform does not support relative " f"constraint {c}" ) - if isinstance(c, ScalarizedOutcomeConstraint): + if len(c.metric_names) > 1: + c_sigs = [ + self._get_metric_signature(n, adapter) for n in c.metric_names + ] # Transform \sum (wi * yi) <= C to # \sum (wi * si * zi) <= C - \sum (wi * mu_i) # Update bound and weights. - c.bound = float( + old_weights = [w for _, w in c.metric_weights] + new_bound = float( c.bound - sum( - c.weights[i] * self.Ymean[(m.signature, strata)] - for i, m in enumerate(c.metrics) + old_weights[i] * self.Ymean[(sig, strata)] + for i, sig in enumerate(c_sigs) ) ) - c.weights = [ - c.weights[i] * self.Ystd[(m.signature, strata)] - for i, m in enumerate(c.metrics) + new_weights = [ + old_weights[i] * self.Ystd[(sig, strata)] + for i, sig in enumerate(c_sigs) ] + new_metric_weights = [ + (name, new_w) + for (name, _), new_w in zip(c.metric_weights, new_weights) + ] + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=new_metric_weights, + op=op_str, + bound=new_bound, + relative=False, + ) + ) + ) else: - c.bound = ( - c.bound - self.Ymean[(c.metric.signature, strata)] - ) / self.Ystd[(c.metric.signature, strata)] - return optimization_config + c_sig = self._get_metric_signature(c.metric_names[0], adapter) + new_bound = float( + (c.bound - self.Ymean[(c_sig, strata)]) / self.Ystd[(c_sig, strata)] + ) + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=op_str, + bound=new_bound, + relative=c.relative, + ) + ) + ) + return new_constraints def untransform_observations( self, @@ -239,32 +307,61 @@ def untransform_outcome_constraints( ) v = none_throws(fixed_features.parameters[self.p_name]) strata = self.strata_mapping[v] + new_constraints = [] for c in outcome_constraints: if c.relative: raise ValueError( "StratifiedStandardizeY does not support relative constraints" ) - if isinstance(c, ScalarizedOutcomeConstraint): + if len(c.metric_names) > 1: + c_sigs = [self._get_metric_signature(n) for n in c.metric_names] # Untransform \sum (wi * si * zi) <= C' back to \sum (wi * yi) <= C # where C' = C - \sum (wi * mu_i) and weights were multiplied by si. # First untransform weights, then untransform bound. - c.weights = [ - c.weights[i] / self.Ystd[(m.signature, strata)] - for i, m in enumerate(c.metrics) + old_weights = [w for _, w in c.metric_weights] + new_weights = [ + old_weights[i] / self.Ystd[(sig, strata)] + for i, sig in enumerate(c_sigs) ] - c.bound = float( + new_bound = float( c.bound + sum( - c.weights[i] * self.Ymean[(m.signature, strata)] - for i, m in enumerate(c.metrics) + new_weights[i] * self.Ymean[(sig, strata)] + for i, sig in enumerate(c_sigs) + ) + ) + new_metric_weights = [ + (name, new_w) + for (name, _), new_w in zip(c.metric_weights, new_weights) + ] + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=new_metric_weights, + op=op_str, + bound=new_bound, + relative=False, + ) ) ) else: - c.bound = float( - c.bound * self.Ystd[(c.metric.signature, strata)] - + self.Ymean[(c.metric.signature, strata)] + c_sig = self._get_metric_signature(c.metric_names[0]) + new_bound = float( + c.bound * self.Ystd[(c_sig, strata)] + self.Ymean[(c_sig, strata)] + ) + op_str = ">=" if c.op == ComparisonOp.GEQ else "<=" + new_constraints.append( + OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=op_str, + bound=new_bound, + relative=c.relative, + ) + ) ) - return outcome_constraints + return new_constraints def transform_experiment_data( self, experiment_data: ExperimentData diff --git a/ax/adapter/transforms/tests/test_add_execution_viability.py b/ax/adapter/transforms/tests/test_add_execution_viability.py index c7a888130cb..5c68cad6756 100644 --- a/ax/adapter/transforms/tests/test_add_execution_viability.py +++ b/ax/adapter/transforms/tests/test_add_execution_viability.py @@ -66,7 +66,7 @@ def test_transform_optimization_config(self) -> None: ) # Check that the last constraint is the feasibility constraint feasibility_constraint = new_opt_config.outcome_constraints[-1] - self.assertEqual(feasibility_constraint.metric.name, "execution_viable") + self.assertEqual(feasibility_constraint.metric_names[0], "execution_viable") self.assertEqual(feasibility_constraint.op, ComparisonOp.GEQ) self.assertEqual(feasibility_constraint.bound, 0.8) self.assertFalse(feasibility_constraint.relative) diff --git a/ax/adapter/transforms/tests/test_base_transform.py b/ax/adapter/transforms/tests/test_base_transform.py index c2e932daeef..f4802aafe5e 100644 --- a/ax/adapter/transforms/tests/test_base_transform.py +++ b/ax/adapter/transforms/tests/test_base_transform.py @@ -190,8 +190,10 @@ def test_transform_optimization_config_preserves_other_fields(self) -> None: ) # Confirm other fields are preserved - self.assertEqual(transformed_config.objective.metric.name, "m1") + self.assertEqual(transformed_config.objective.metric_names[0], "m1") self.assertTrue(transformed_config.objective.minimize) self.assertEqual(len(transformed_config.outcome_constraints), 1) - self.assertEqual(transformed_config.outcome_constraints[0].metric.name, "m2") + self.assertEqual( + transformed_config.outcome_constraints[0].metric_names[0], "m2" + ) self.assertEqual(transformed_config.outcome_constraints[0].bound, 10.0) diff --git a/ax/adapter/transforms/tests/test_derelativize_transform.py b/ax/adapter/transforms/tests/test_derelativize_transform.py index 228f29c8e2f..e756d8a31ec 100644 --- a/ax/adapter/transforms/tests/test_derelativize_transform.py +++ b/ax/adapter/transforms/tests/test_derelativize_transform.py @@ -43,6 +43,7 @@ def test_DerelativizeTransform(self) -> None: for negative_metrics in [False, True]: sq_sign = -1.0 if negative_metrics else 1.0 observed_means = sq_sign * np.array([[1.0, 2.0], [1.0, 6.0]]) + sq_a_observed = float(np.mean(observed_means[1, 0])) sq_b_observed = np.mean(observed_means[1, 1]) experiment = get_experiment_with_observations( observations=observed_means.tolist(), @@ -51,6 +52,7 @@ def test_DerelativizeTransform(self) -> None: ) predicted_means = sq_sign * np.array([3.0, 5.0]) + sq_a_predicted = predicted_means[0] sq_b_predicted = predicted_means[-1] predict_return_value = [ ObservationData( @@ -70,7 +72,9 @@ def test_DerelativizeTransform(self) -> None: self._test_DerelativizeTransform( experiment=experiment, mock_predict=mock_predict, + sq_a_observed=sq_a_observed, sq_b_observed=sq_b_observed, + sq_a_predicted=sq_a_predicted, sq_b_predicted=sq_b_predicted, ) @@ -78,7 +82,9 @@ def _test_DerelativizeTransform( self, experiment: Experiment, mock_predict: Mock, + sq_a_observed: float, sq_b_observed: float, + sq_a_predicted: float, sq_b_predicted: float, ) -> None: t = Derelativize(search_space=None) @@ -98,12 +104,12 @@ def _test_DerelativizeTransform( g = Adapter(experiment=experiment_1, generator=Generator()) # Test with no relative constraints - objective = Objective(Metric("c"), minimize=True) + objective = Objective(metric=Metric("c"), minimize=True) oc = OptimizationConfig( objective=objective, outcome_constraints=[ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=False ), ScalarizedOutcomeConstraint( metrics=[Metric("m1"), Metric("m2")], @@ -123,14 +129,17 @@ def _test_DerelativizeTransform( objective=objective, outcome_constraints=[ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=False ), OutcomeConstraint( - Metric("m2"), ComparisonOp.LEQ, bound=relative_bound, relative=True + metric=Metric("m2"), + op=ComparisonOp.LEQ, + bound=relative_bound, + relative=True, ), ScalarizedOutcomeConstraint( metrics=[Metric("m1"), Metric("m2")], - weights=[0.0, 1.0], + weights=[0.5, 0.5], op=ComparisonOp.LEQ, bound=relative_bound, relative=True, @@ -140,20 +149,27 @@ def _test_DerelativizeTransform( oc2 = t.transform_optimization_config(oc, g, None) sq_val = sq_b_predicted absolute_bound = (1 + np.sign(sq_val) * relative_bound / 100.0) * sq_val + sq_val_scalarized = 0.5 * sq_a_predicted + 0.5 * sq_b_predicted + absolute_bound_scalarized = ( + 1 + np.sign(sq_val_scalarized) * relative_bound / 100.0 + ) * sq_val_scalarized self.assertTrue( oc2.outcome_constraints == [ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=False ), OutcomeConstraint( - Metric("m2"), ComparisonOp.LEQ, bound=absolute_bound, relative=False + metric=Metric("m2"), + op=ComparisonOp.LEQ, + bound=absolute_bound, + relative=False, ), ScalarizedOutcomeConstraint( metrics=[Metric("m1"), Metric("m2")], - weights=[0.0, 1.0], + weights=[0.5, 0.5], op=ComparisonOp.LEQ, - bound=absolute_bound, + bound=absolute_bound_scalarized, relative=False, ), ] @@ -179,14 +195,17 @@ def _test_DerelativizeTransform( objective=objective, outcome_constraints=[ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=False ), OutcomeConstraint( - Metric("m2"), ComparisonOp.LEQ, bound=relative_bound, relative=True + metric=Metric("m2"), + op=ComparisonOp.LEQ, + bound=relative_bound, + relative=True, ), ScalarizedOutcomeConstraint( metrics=[Metric("m1"), Metric("m2")], - weights=[0.0, 1.0], + weights=[0.5, 0.5], op=ComparisonOp.LEQ, bound=relative_bound, relative=True, @@ -196,20 +215,27 @@ def _test_DerelativizeTransform( oc2 = t.transform_optimization_config(oc, g, None) sq_val = sq_b_observed absolute_bound = (1 + np.sign(sq_val) * relative_bound / 100.0) * sq_val + sq_val_scalarized = 0.5 * sq_a_observed + 0.5 * sq_b_observed + absolute_bound_scalarized = ( + 1 + np.sign(sq_val_scalarized) * relative_bound / 100.0 + ) * sq_val_scalarized self.assertEqual( oc2.outcome_constraints, [ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=False ), OutcomeConstraint( - Metric("m2"), ComparisonOp.LEQ, bound=absolute_bound, relative=False + metric=Metric("m2"), + op=ComparisonOp.LEQ, + bound=absolute_bound, + relative=False, ), ScalarizedOutcomeConstraint( metrics=[Metric("m1"), Metric("m2")], - weights=[0.0, 1.0], + weights=[0.5, 0.5], op=ComparisonOp.LEQ, - bound=absolute_bound, + bound=absolute_bound_scalarized, relative=False, ), ], @@ -227,10 +253,10 @@ def _test_DerelativizeTransform( objective=objective, outcome_constraints=[ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=False ), OutcomeConstraint( - Metric("m2"), ComparisonOp.LEQ, bound=-10, relative=True + metric=Metric("m2"), op=ComparisonOp.LEQ, bound=-10, relative=True ), ], ) @@ -293,10 +319,10 @@ def _test_DerelativizeTransform( def test_errors(self) -> None: t = Derelativize(search_space=None) oc = OptimizationConfig( - objective=Objective(Metric("c"), minimize=False), + objective=Objective(metric=Metric("c"), minimize=False), outcome_constraints=[ OutcomeConstraint( - Metric("m1"), ComparisonOp.LEQ, bound=2, relative=True + metric=Metric("m1"), op=ComparisonOp.LEQ, bound=2, relative=True ) ], ) diff --git a/ax/adapter/transforms/tests/test_log_y_transform.py b/ax/adapter/transforms/tests/test_log_y_transform.py index f2982fba108..21482edce76 100644 --- a/ax/adapter/transforms/tests/test_log_y_transform.py +++ b/ax/adapter/transforms/tests/test_log_y_transform.py @@ -28,7 +28,7 @@ MultiObjectiveOptimizationConfig, OptimizationConfig, ) -from ax.core.outcome_constraint import ObjectiveThreshold +from ax.core.outcome_constraint import ObjectiveThreshold, OutcomeConstraint from ax.generators.torch.botorch_modular.generator import BoTorchGenerator from ax.generators.types import TConfig from ax.utils.common.testutils import TestCase @@ -191,7 +191,12 @@ def test_TransformOptimizationConfig(self) -> None: ], ) oc_tf = tf.transform_optimization_config(deepcopy(oc), None, None) - oc.outcome_constraints[0].bound = math.log(2.345) + oc.outcome_constraints[0] = OutcomeConstraint( + metric=m1, + op=oc.outcome_constraints[0].op, + bound=math.log(2.345), + relative=oc.outcome_constraints[0].relative, + ) self.assertEqual(oc_tf, oc) # output constraint with a relative bound should fail oc = OptimizationConfig( @@ -218,7 +223,7 @@ def test_TransformOptimizationConfigMOO(self) -> None: Objective(metric=m2, minimize=True), ], ) - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] = [ ObjectiveThreshold(metric=m1, bound=1.234, relative=False), ObjectiveThreshold(metric=m2, bound=3.456, relative=False), ] @@ -228,7 +233,9 @@ def test_TransformOptimizationConfigMOO(self) -> None: ) tf = LogY(search_space=None, config={"metrics": ["m1"]}) oc_tf = tf.transform_optimization_config(deepcopy(oc), None, None) - oc.objective_thresholds[0].bound = math.log(1.234) + oc.objective_thresholds[0] = ObjectiveThreshold( + metric=m1, bound=math.log(1.234), relative=False + ) self.assertEqual(oc_tf, oc) def _base_test_transform_experiment_data( diff --git a/ax/adapter/transforms/tests/test_objective_as_constraint.py b/ax/adapter/transforms/tests/test_objective_as_constraint.py index e658ae4e9de..37a70dddc93 100644 --- a/ax/adapter/transforms/tests/test_objective_as_constraint.py +++ b/ax/adapter/transforms/tests/test_objective_as_constraint.py @@ -151,7 +151,7 @@ def test_adds_constraint_when_no_feasible_points(self) -> None: self.assertEqual(len(transformed.outcome_constraints), 2) new_constraint = transformed.outcome_constraints[1] - self.assertEqual(new_constraint.metric.name, "m1") + self.assertEqual(new_constraint.metric_names[0], "m1") self.assertEqual(new_constraint.op, ComparisonOp.GEQ) self.assertEqual(new_constraint.bound, 1.0) # SQ value for m1 self.assertFalse(new_constraint.relative) @@ -179,7 +179,7 @@ def test_adds_leq_constraint_when_minimizing(self) -> None: transformed = t.transform_optimization_config(opt_config, adapter) new_constraint = transformed.outcome_constraints[1] - self.assertEqual(new_constraint.metric.name, "m1") + self.assertEqual(new_constraint.metric_names[0], "m1") self.assertEqual(new_constraint.op, ComparisonOp.LEQ) self.assertEqual(new_constraint.bound, 1.0) # SQ value for m1 self.assertFalse(new_constraint.relative) @@ -187,11 +187,13 @@ def test_adds_leq_constraint_when_minimizing(self) -> None: def test_no_op_without_status_quo(self) -> None: """Test that the transform is a no-op without a status quo.""" optimization_config = OptimizationConfig( - objective=Objective(Metric("m1", lower_is_better=False), minimize=False), + objective=Objective( + metric=Metric("m1", lower_is_better=False), minimize=False + ), outcome_constraints=[ OutcomeConstraint( - Metric("m2", lower_is_better=True), - ComparisonOp.GEQ, + metric=Metric("m2", lower_is_better=True), + op=ComparisonOp.GEQ, bound=10.0, ), ], @@ -218,7 +220,9 @@ def test_no_op_without_status_quo(self) -> None: def test_no_op_without_constraints(self) -> None: """Test that the transform is a no-op when there are no constraints.""" optimization_config = OptimizationConfig( - objective=Objective(Metric("m1", lower_is_better=False), minimize=False), + objective=Objective( + metric=Metric("m1", lower_is_better=False), minimize=False + ), ) search_space = SearchSpace( parameters=[ @@ -273,7 +277,7 @@ def test_untransform_removes_added_constraint(self) -> None: outcome_constraints=transformed.outcome_constraints, ) self.assertEqual(len(untransformed), 1) - self.assertEqual(untransformed[0].metric.name, "m2") + self.assertEqual(untransformed[0].metric_names[0], "m2") def test_raises_on_relative_constraints(self) -> None: """Test that a ValueError is raised in transform_optimization_config @@ -500,10 +504,10 @@ def test_scalarized_objective_soo_adds_scalarized_constraint(self) -> None: ] self.assertEqual(len(scalarized_constraints), 1) sc = scalarized_constraints[0] + # The expression "m1 + m2 >= 3" round-trips as GEQ with bound 3. self.assertEqual(sc.op, ComparisonOp.GEQ) self.assertAlmostEqual(sc.bound, 3.0) - self.assertEqual([m.name for m in sc.metrics], ["m1", "m2"]) - self.assertEqual(sc.weights, [1.0, 1.0]) + self.assertCountEqual(sc.metric_names, ["m1", "m2"]) self.assertFalse(sc.relative) def test_scalarized_objective_soo_minimize(self) -> None: @@ -534,8 +538,9 @@ def test_scalarized_objective_soo_minimize(self) -> None: ] self.assertEqual(len(scalarized_constraints), 1) sc = scalarized_constraints[0] + # With sign-encoded weights [-1, -1], the expression is + # "-m1 - m2 >= -3" which normalizes to LEQ form: "m1 + m2 <= 3". self.assertEqual(sc.op, ComparisonOp.LEQ) - # SQ value = 1.0*1.0 + 1.0*2.0 = 3.0 self.assertAlmostEqual(sc.bound, 3.0) def test_scalarized_objective_with_weights(self) -> None: @@ -569,9 +574,9 @@ def test_scalarized_objective_with_weights(self) -> None: ] self.assertEqual(len(scalarized_constraints), 1) sc = scalarized_constraints[0] - self.assertEqual(sc.op, ComparisonOp.GEQ) + # Expression "2*m1 - m2 >= 0" normalizes to LEQ: "-2*m1 + m2 <= 0". + self.assertEqual(sc.op, ComparisonOp.LEQ) self.assertAlmostEqual(sc.bound, 0.0) - self.assertEqual(sc.weights, [2.0, -1.0]) def test_scalarized_objective_moo_config(self) -> None: """Test ScalarizedObjective in MultiObjectiveOptimizationConfig.""" @@ -604,6 +609,7 @@ def test_scalarized_objective_moo_config(self) -> None: ] self.assertEqual(len(scalarized_constraints), 1) sc = scalarized_constraints[0] + # Expression "m1 + m2 >= 3" round-trips as GEQ with bound 3. self.assertEqual(sc.op, ComparisonOp.GEQ) self.assertAlmostEqual(sc.bound, 3.0) @@ -631,7 +637,7 @@ def test_scalarized_objective_untransform(self) -> None: outcome_constraints=transformed.outcome_constraints, ) self.assertEqual(len(untransformed), 1) - self.assertEqual(untransformed[0].metric.name, "m3") + self.assertEqual(untransformed[0].metric_names[0], "m3") def test_scalarized_objective_no_op_when_feasible(self) -> None: """Test no-op with ScalarizedObjective when feasible points exist.""" @@ -703,7 +709,7 @@ def _make_moo_experiment_adapter_and_data( ) ) - obj_thresholds_list: list[ObjectiveThreshold] = [] + obj_thresholds_list: list[OutcomeConstraint] = [] if objective_thresholds is not None: for i, (bound, relative) in enumerate(objective_thresholds): metric_name = f"m{i + 1}" @@ -821,14 +827,14 @@ def test_moo_adds_constraints_when_no_feasible_points(self) -> None: # Check the new constraints. m1_constraint = next( - c for c in transformed.outcome_constraints if c.metric.name == "m1" + c for c in transformed.outcome_constraints if c.metric_names[0] == "m1" ) self.assertEqual(m1_constraint.op, ComparisonOp.GEQ) self.assertEqual(m1_constraint.bound, 1.0) self.assertFalse(m1_constraint.relative) m2_constraint = next( - c for c in transformed.outcome_constraints if c.metric.name == "m2" + c for c in transformed.outcome_constraints if c.metric_names[0] == "m2" ) self.assertEqual(m2_constraint.op, ComparisonOp.GEQ) self.assertEqual(m2_constraint.bound, 1.0) @@ -858,13 +864,13 @@ def test_moo_adds_leq_constraints_when_minimizing(self) -> None: transformed = t.transform_optimization_config(opt_config, adapter) m1_constraint = next( - c for c in transformed.outcome_constraints if c.metric.name == "m1" + c for c in transformed.outcome_constraints if c.metric_names[0] == "m1" ) self.assertEqual(m1_constraint.op, ComparisonOp.LEQ) self.assertEqual(m1_constraint.bound, 1.0) m2_constraint = next( - c for c in transformed.outcome_constraints if c.metric.name == "m2" + c for c in transformed.outcome_constraints if c.metric_names[0] == "m2" ) self.assertEqual(m2_constraint.op, ComparisonOp.LEQ) self.assertEqual(m2_constraint.bound, 2.0) @@ -912,4 +918,4 @@ def test_moo_untransform_removes_added_constraints(self) -> None: outcome_constraints=transformed.outcome_constraints, ) self.assertEqual(len(untransformed), 1) - self.assertEqual(untransformed[0].metric.name, "m3") + self.assertEqual(untransformed[0].metric_names[0], "m3") diff --git a/ax/adapter/transforms/tests/test_power_y_transform.py b/ax/adapter/transforms/tests/test_power_y_transform.py index 04abe99bbb0..737e8718567 100644 --- a/ax/adapter/transforms/tests/test_power_y_transform.py +++ b/ax/adapter/transforms/tests/test_power_y_transform.py @@ -240,7 +240,12 @@ def test_transform_optimization_config(self) -> None: tf_bound = ( tf.power_transforms["m1"].transform(np.array(bound, ndmin=2)).item() ) - oc_true.outcome_constraints[0].bound = tf_bound + oc_true.outcome_constraints[0] = OutcomeConstraint( + metric=m1, + op=oc_true.outcome_constraints[0].op, + bound=tf_bound, + relative=oc_true.outcome_constraints[0].relative, + ) self.assertEqual(oc_tf, oc_true) # Check untransform of outcome constraint cons = tf.untransform_outcome_constraints( diff --git a/ax/adapter/transforms/tests/test_relativize_transform.py b/ax/adapter/transforms/tests/test_relativize_transform.py index 2e1431a81b9..a4ce889f3dd 100644 --- a/ax/adapter/transforms/tests/test_relativize_transform.py +++ b/ax/adapter/transforms/tests/test_relativize_transform.py @@ -26,6 +26,7 @@ from ax.core import BatchTrial from ax.core.arm import Arm from ax.core.experiment import Experiment +from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.observation_utils import observations_from_data, recombine_observations @@ -544,8 +545,15 @@ def test_transform_optimization_config_with_relative_thresholds(self) -> None: optimization_config = get_branin_multi_objective_optimization_config( has_objective_thresholds=True, ) - for threshold in optimization_config.objective_thresholds: - threshold.relative = True + optimization_config.objective_thresholds = [ + OutcomeConstraint( + metric=Metric(name=threshold.metric_names[0]), + op=threshold.op, + bound=threshold.bound, + relative=True, + ) + for threshold in optimization_config.objective_thresholds + ] new_config = relativize.transform_optimization_config( optimization_config=optimization_config, @@ -575,7 +583,13 @@ def test_transform_optimization_config_with_non_relative_thresholds(self) -> Non optimization_config = get_branin_multi_objective_optimization_config( has_objective_thresholds=True, ) - optimization_config.objective_thresholds[1].relative = False + old_threshold = optimization_config.objective_thresholds[1] + optimization_config.objective_thresholds[1] = OutcomeConstraint( + metric=Metric(name=old_threshold.metric_names[0]), + op=old_threshold.op, + bound=old_threshold.bound, + relative=False, + ) with self.assertRaisesRegex( ValueError, "All objective thresholds must be relative" @@ -624,6 +638,7 @@ def make_pref_config(expect_relativized: bool) -> PreferenceOptimizationConfig: Objective( metric=BraninMetric("branin", ["x1", "x2"]), minimize=False ), + Objective(metric=Metric(name="aux_metric"), minimize=False), ] ), preference_profile_name="test", diff --git a/ax/adapter/transforms/tests/test_standardize_y_transform.py b/ax/adapter/transforms/tests/test_standardize_y_transform.py index 89bc4b68cba..6985b589328 100644 --- a/ax/adapter/transforms/tests/test_standardize_y_transform.py +++ b/ax/adapter/transforms/tests/test_standardize_y_transform.py @@ -24,7 +24,6 @@ from ax.utils.testing.core_stubs import get_experiment_with_observations from pandas import DataFrame from pandas.testing import assert_frame_equal -from pyre_extensions import assert_is_instance class StandardizeYTransformTest(TestCase): @@ -148,25 +147,27 @@ def test_TransformOptimizationConfig(self) -> None: ] oc = OptimizationConfig(objective=objective, outcome_constraints=cons) oc = self.t.transform_optimization_config(oc, None, None) - cons_t = [ - OutcomeConstraint( - metric=m1, op=ComparisonOp.GEQ, bound=1.0, relative=False - ), - OutcomeConstraint( - metric=m2, - op=ComparisonOp.LEQ, - bound=2.0 * sqrt(3), # (3.5 - 1.5) / sqrt(1/3) - relative=False, - ), - ScalarizedOutcomeConstraint( - metrics=[m1, m2], - weights=[0.5 * 1.0, 0.5 * sqrt(1 / 3)], - op=ComparisonOp.LEQ, - bound=2.25, # 3.5 - (0.5 * 1.0 + 0.5 * 1.5) - relative=False, - ), - ] - self.assertTrue(oc.outcome_constraints == cons_t) + # Verify the transformed constraints have the expected values. + # We compare properties individually to avoid floating-point string + # representation issues with expression-based equality. + self.assertEqual(len(oc.outcome_constraints), 3) + c0, c1, c2 = oc.outcome_constraints + # Constraint 0: m1 >= (2.0 - 1.0) / 1.0 = 1.0 + self.assertEqual(c0.op, ComparisonOp.GEQ) + self.assertAlmostEqual(c0.bound, 1.0) + self.assertFalse(c0.relative) + # Constraint 1: m2 <= (3.5 - 1.5) / sqrt(1/3) = 2*sqrt(3) + self.assertEqual(c1.op, ComparisonOp.LEQ) + self.assertAlmostEqual(c1.bound, 2.0 * sqrt(3)) + self.assertFalse(c1.relative) + # Constraint 2 (scalarized): 0.5*m1 + 0.5*sqrt(1/3)*m2 <= 2.25 + self.assertEqual(c2.op, ComparisonOp.LEQ) + self.assertAlmostEqual(c2.bound, 2.25) + self.assertFalse(c2.relative) + self.assertEqual(c2.metric_names, ["m1", "m2"]) + weights = [w for _, w in c2.metric_weights] + self.assertAlmostEqual(weights[0], 0.5 * 1.0) + self.assertAlmostEqual(weights[1], 0.5 * sqrt(1 / 3)) self.assertTrue(oc.objective == objective) # Check fail with relative @@ -233,10 +234,12 @@ def test_TransformOptimizationConfigWithScalarizedObjective(self) -> None: # Check that weights are scaled by standard deviations expected_weights = [0.5 * 1.0, 0.5 * sqrt(1 / 3)] - transformed_objective = assert_is_instance( - oc_transformed.objective, ScalarizedObjective + transformed_objective = oc_transformed.objective + self.assertTrue( + np.allclose( + [w for _, w in transformed_objective.metric_weights], expected_weights + ) ) - self.assertTrue(np.allclose(transformed_objective.weights, expected_weights)) # Test with ScalarizedObjective missing a metric objective_missing = ScalarizedObjective( @@ -257,14 +260,15 @@ def test_TransformOptimizationConfigWithScalarizedObjective(self) -> None: oc_minimize, None, None ) - # Check that weights are scaled by standard deviations - expected_weights_minimize = [1.0 * 1.0, -2.0 * sqrt(1 / 3)] - transformed_objective_minimize = assert_is_instance( - oc_minimize_transformed.objective, ScalarizedObjective - ) + # Check that weights are scaled by standard deviations. + # Note: ScalarizedObjective with minimize=True bakes negation into the + # expression, so the stored weights are [-1.0, 2.0] (not [1.0, -2.0]). + expected_weights_minimize = [-1.0 * 1.0, 2.0 * sqrt(1 / 3)] + transformed_objective_minimize = oc_minimize_transformed.objective self.assertTrue( np.allclose( - transformed_objective_minimize.weights, expected_weights_minimize + [w for _, w in transformed_objective_minimize.metric_weights], + expected_weights_minimize, ) ) diff --git a/ax/adapter/transforms/tests/test_stratified_standardize_y_transform.py b/ax/adapter/transforms/tests/test_stratified_standardize_y_transform.py index 1da962ff4eb..52b1f1169bc 100644 --- a/ax/adapter/transforms/tests/test_stratified_standardize_y_transform.py +++ b/ax/adapter/transforms/tests/test_stratified_standardize_y_transform.py @@ -27,7 +27,6 @@ from ax.utils.testing.core_stubs import get_experiment_with_observations from pandas import DataFrame from pandas.testing import assert_frame_equal -from pyre_extensions import assert_is_instance class StratifiedStandardizeYTransformTest(TestCase): @@ -318,18 +317,16 @@ def test_TransformOptimizationConfig(self) -> None: oc = OptimizationConfig(objective=self.objective, outcome_constraints=cons2) fixed_features = ObservationFeatures({"z": "a"}) oc = self.t.transform_optimization_config(oc, None, fixed_features) - cons_t = [ - OutcomeConstraint( - metric=self.m1, op=ComparisonOp.GEQ, bound=1.0, relative=False - ), - OutcomeConstraint( - metric=self.m2, - op=ComparisonOp.LEQ, - bound=(3.5 - 5.0) / (sqrt(2) * 3), - relative=False, - ), - ] - self.assertTrue(oc.outcome_constraints == cons_t) + # Verify constraint values approximately (expression-string equality + # can fail due to floating-point formatting differences). + self.assertEqual(len(oc.outcome_constraints), 2) + c0, c1 = oc.outcome_constraints + self.assertEqual(c0.op, ComparisonOp.GEQ) + self.assertAlmostEqual(c0.bound, 1.0) + self.assertFalse(c0.relative) + self.assertEqual(c1.op, ComparisonOp.LEQ) + self.assertAlmostEqual(c1.bound, (3.5 - 5.0) / (sqrt(2) * 3)) + self.assertFalse(c1.relative) self.assertTrue(oc.objective == self.objective) # No constraints @@ -355,53 +352,39 @@ def test_TransformOptimizationConfigWithStrataMapping(self) -> None: cons2 = deepcopy(self.cons) oc = OptimizationConfig(objective=self.objective, outcome_constraints=cons2) fixed_features = ObservationFeatures({"z": "a"}) - cons_t = [ - OutcomeConstraint( - metric=self.m1, op=ComparisonOp.GEQ, bound=1.0, relative=False - ), - OutcomeConstraint( - metric=self.m2, - op=ComparisonOp.LEQ, - bound=(3.5 - 5.0) / (sqrt(2) * 3), - relative=False, - ), - ] cons2 = deepcopy(self.cons) oc = OptimizationConfig(objective=self.objective, outcome_constraints=cons2) oc = self.t2.transform_optimization_config(oc, None, fixed_features) - self.assertTrue(oc.outcome_constraints == cons_t) + # Verify constraint values approximately. + self.assertEqual(len(oc.outcome_constraints), 2) + c0, c1 = oc.outcome_constraints + self.assertEqual(c0.op, ComparisonOp.GEQ) + self.assertAlmostEqual(c0.bound, 1.0) + self.assertFalse(c0.relative) + self.assertEqual(c1.op, ComparisonOp.LEQ) + self.assertAlmostEqual(c1.bound, (3.5 - 5.0) / (sqrt(2) * 3)) + self.assertFalse(c1.relative) self.assertTrue(oc.objective == self.objective) fixed_features = ObservationFeatures({"z": "c"}) cons2 = deepcopy(self.cons) oc = OptimizationConfig(objective=self.objective, outcome_constraints=cons2) oc = self.t2.transform_optimization_config(oc, None, fixed_features) - cons_t2 = [ - OutcomeConstraint( - metric=self.m1, - op=ComparisonOp.GEQ, - bound=(2.0 - self.t2.Ymean[("m1", 1)]) / self.t2.Ystd[("m1", 1)], - relative=False, - ), - OutcomeConstraint( - metric=self.m2, - op=ComparisonOp.LEQ, - bound=(3.5 - self.t2.Ymean[("m2", 1)]) / self.t2.Ystd[("m2", 1)], - relative=False, - ), - ] - self.assertEqual(oc.outcome_constraints, cons_t2) + # Verify constraint values approximately. + self.assertEqual(len(oc.outcome_constraints), 2) + c0, c1 = oc.outcome_constraints + self.assertEqual(c0.op, ComparisonOp.GEQ) + self.assertAlmostEqual( + c0.bound, + (2.0 - self.t2.Ymean[("m1", 1)]) / self.t2.Ystd[("m1", 1)], + ) + self.assertFalse(c0.relative) + self.assertEqual(c1.op, ComparisonOp.LEQ) + self.assertAlmostEqual( + c1.bound, + (3.5 - self.t2.Ymean[("m2", 1)]) / self.t2.Ystd[("m2", 1)], + ) + self.assertFalse(c1.relative) self.assertTrue(oc.objective == self.objective) - cons_t = [ - OutcomeConstraint( - metric=self.m1, op=ComparisonOp.GEQ, bound=1.0, relative=False - ), - OutcomeConstraint( - metric=self.m2, - op=ComparisonOp.LEQ, - bound=(3.5 - 5.0) / (sqrt(2) * 3), - relative=False, - ), - ] def test_transform_experiment_data(self) -> None: transformed_data = self.t.transform_experiment_data( @@ -470,11 +453,12 @@ def test_TransformOptimizationConfigWithScalarizedObjective(self) -> None: adapter=None, fixed_features=fixed_features, ) - transformed_objective = assert_is_instance( - oc_transformed.objective, ScalarizedObjective - ) + transformed_objective = oc_transformed.objective self.assertTrue( - np.allclose(transformed_objective.weights, expected_weights[strata]) + np.allclose( + [w for _, w in transformed_objective.metric_weights], + expected_weights[strata], + ) ) def test_TransformAndUntransformScalarizedOutcomeConstraint(self) -> None: @@ -500,20 +484,25 @@ def test_TransformAndUntransformScalarizedOutcomeConstraint(self) -> None: # For strata "a": mu_m1 = 1.0, mu_m2 = 5.0, std_m1 = 1.0, std_m2 = sqrt(2) * 3 # new_bound = 10.0 - (1.0 * 1.0 + 2.0 * 5.0) = 10.0 - 11.0 = -1.0 # new_weights = [1.0 * 1.0, 2.0 * sqrt(2) * 3] - transformed_constraint = assert_is_instance( - oc_transformed.outcome_constraints[0], ScalarizedOutcomeConstraint - ) + transformed_constraint = oc_transformed.outcome_constraints[0] expected_bound = 10.0 - (1.0 * 1.0 + 2.0 * 5.0) expected_weights = [1.0 * 1.0, 2.0 * sqrt(2) * 3] self.assertAlmostEqual(transformed_constraint.bound, expected_bound) - self.assertTrue(np.allclose(transformed_constraint.weights, expected_weights)) + self.assertTrue( + np.allclose( + [w for _, w in transformed_constraint.metric_weights], expected_weights + ) + ) # Untransform and verify round-trip recovers original values untransformed_constraints = self.t.untransform_outcome_constraints( oc_transformed.outcome_constraints, fixed_features ) - untransformed_constraint = assert_is_instance( - untransformed_constraints[0], ScalarizedOutcomeConstraint - ) + untransformed_constraint = untransformed_constraints[0] self.assertAlmostEqual(untransformed_constraint.bound, original_bound) - self.assertTrue(np.allclose(untransformed_constraint.weights, original_weights)) + self.assertTrue( + np.allclose( + [w for _, w in untransformed_constraint.metric_weights], + original_weights, + ) + ) diff --git a/ax/adapter/transforms/tests/test_winsorize_transform.py b/ax/adapter/transforms/tests/test_winsorize_transform.py index 086bd1d3701..6481b0f256e 100644 --- a/ax/adapter/transforms/tests/test_winsorize_transform.py +++ b/ax/adapter/transforms/tests/test_winsorize_transform.py @@ -216,7 +216,12 @@ def test_winsorize_outcome_constraints(self) -> None: ) self.assertEqual(cutoffs, (-INF, 23.5)) # From above with a tight bound - outcome_constraint_leq.bound = 2 + outcome_constraint_leq = OutcomeConstraint( + metric=ma, + op=outcome_constraint_leq.op, + bound=2, + relative=outcome_constraint_leq.relative, + ) cutoffs = _get_auto_winsorization_cutoffs_outcome_constraint( metric_values=self.values, outcome_constraints=[outcome_constraint_leq] ) @@ -227,13 +232,17 @@ def test_winsorize_outcome_constraints(self) -> None: ) self.assertEqual(cutoffs, (-31.5, INF)) # From below with a tight bound - outcome_constraint_geq.bound = 5 + outcome_constraint_geq = OutcomeConstraint( + metric=mb, + op=outcome_constraint_geq.op, + bound=5, + relative=outcome_constraint_geq.relative, + ) cutoffs = _get_auto_winsorization_cutoffs_outcome_constraint( metric_values=self.values, outcome_constraints=[outcome_constraint_geq] ) self.assertEqual(cutoffs, (-6.5, INF)) # Both with the tight bounds - outcome_constraint_geq.bound = 5 cutoffs = _get_auto_winsorization_cutoffs_outcome_constraint( metric_values=self.values, outcome_constraints=[outcome_constraint_leq, outcome_constraint_geq], @@ -391,13 +400,29 @@ def test_winsorization_with_optimization_config(self) -> None: self.assertEqual(transform.cutoffs["m2"], (-INF, 10.0)) # 4 + 1.5 * 4 self.assertEqual(transform.cutoffs["m3"], (-INF, INF)) # Make the constraint absolute, which should trigger winsorization - outcome_constraint.relative = False + outcome_constraint = OutcomeConstraint( + metric=m1, + op=outcome_constraint.op, + bound=outcome_constraint.bound, + relative=False, + ) + none_throws(experiment.optimization_config).outcome_constraints = [ + outcome_constraint + ] transform = Winsorize(experiment_data=experiment_data, adapter=adapter) self.assertEqual(transform.cutoffs["m1"], (-INF, 13.5)) # 6 + 1.5 * 5 self.assertEqual(transform.cutoffs["m2"], (-INF, 10.0)) # 4 + 1.5 * 4 self.assertEqual(transform.cutoffs["m3"], (-INF, INF)) # Change to a GEQ constraint - outcome_constraint.op = ComparisonOp.GEQ + outcome_constraint = OutcomeConstraint( + metric=m1, + op=ComparisonOp.GEQ, + bound=outcome_constraint.bound, + relative=outcome_constraint.relative, + ) + none_throws(experiment.optimization_config).outcome_constraints = [ + outcome_constraint + ] transform = Winsorize(experiment_data=experiment_data, adapter=adapter) self.assertEqual(transform.cutoffs["m1"], (-6.5, INF)) # 1 - 1.5 * 5 self.assertEqual(transform.cutoffs["m2"], (-INF, 10.0)) # 4 + 1.5 * 4 @@ -419,7 +444,7 @@ def test_winsorization_with_optimization_config(self) -> None: ) # Multi-objective without objective thresholds should warn and winsorize moo_objective = MultiObjective( - [Objective(m1, minimize=False), Objective(m2, minimize=True)] + [Objective(metric=m1, minimize=False), Objective(metric=m2, minimize=True)] ) optimization_config = MultiObjectiveOptimizationConfig(objective=moo_objective) experiment._optimization_config = optimization_config @@ -437,7 +462,7 @@ def test_winsorization_with_optimization_config(self) -> None: self.assertEqual(transform.cutoffs["m2"], (-INF, 10.0)) self.assertEqual(transform.cutoffs["m3"], (-INF, INF)) # Add relative objective thresholds. Should warn and skip. - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] = [ ObjectiveThreshold(m1, 3, relative=True), ObjectiveThreshold(m2, 4, relative=True), ] @@ -458,8 +483,10 @@ def test_winsorization_with_optimization_config(self) -> None: for i in range(1, 4): self.assertEqual(transform.cutoffs[f"m{i}"], (-INF, INF)) # Make the objective thresholds absolute (should trigger winsorization) - optimization_config.objective_thresholds[0].relative = False - optimization_config.objective_thresholds[1].relative = False + optimization_config.objective_thresholds = [ + ObjectiveThreshold(m1, 3, relative=False), + ObjectiveThreshold(m2, 4, relative=False), + ] transform = Winsorize(experiment_data=experiment_data, adapter=adapter) self.assertEqual(transform.cutoffs["m1"], (-6.5, INF)) # 1 - 1.5 * 5 self.assertEqual(transform.cutoffs["m2"], (-INF, 10.0)) # 4 + 1.5 * 4 @@ -483,17 +510,17 @@ def test_relative_constraints(self) -> None: ) # Test with relative constraint, in-design status quo oc = OptimizationConfig( - objective=Objective(Metric("c"), minimize=False), + objective=Objective(metric=Metric("c"), minimize=False), outcome_constraints=[ OutcomeConstraint( - Metric("a"), ComparisonOp.LEQ, bound=2, relative=False + metric=Metric("a"), op=ComparisonOp.LEQ, bound=2, relative=False ), OutcomeConstraint( - Metric("b"), ComparisonOp.LEQ, bound=-10, relative=True + metric=Metric("b"), op=ComparisonOp.LEQ, bound=-10, relative=True ), ScalarizedOutcomeConstraint( metrics=[Metric("a"), Metric("b")], - weights=[0.0, 1.0], + weights=[0.5, 0.5], op=ComparisonOp.LEQ, bound=-10, relative=True, @@ -569,7 +596,7 @@ def test_relative_constraints(self) -> None: config={"derelativize_with_raw_status_quo": True}, ) self.assertDictEqual( - t.cutoffs, {"a": (-INF, 4.25), "b": (-INF, 4.25), "c": (-3.25, INF)} + t.cutoffs, {"a": (-INF, 2.625), "b": (-INF, 4.25), "c": (-3.25, INF)} ) def test_transform_experiment_data(self) -> None: diff --git a/ax/adapter/transforms/winsorize.py b/ax/adapter/transforms/winsorize.py index 441a715380e..c7d69e14143 100644 --- a/ax/adapter/transforms/winsorize.py +++ b/ax/adapter/transforms/winsorize.py @@ -18,7 +18,7 @@ from ax.adapter.transforms.utils import ( derelativize_optimization_config_with_raw_status_quo, ) -from ax.core.objective import MultiObjective, ScalarizedObjective +from ax.core.objective import ScalarizedObjective from ax.core.observation import ObservationData from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, @@ -214,11 +214,9 @@ def _get_cutoffs( if adapter is None or optimization_config is None: return DEFAULT_CUTOFFS relative_constraint_metrics = { - metric.name + name for oc in optimization_config.all_constraints - for metric in ( - oc.metrics if isinstance(oc, ScalarizedOutcomeConstraint) else [oc.metric] - ) + for name in oc.metric_names if oc.relative } if metric_signature in relative_constraint_metrics: @@ -237,7 +235,7 @@ def _get_cutoffs( ) # Non-objective metrics - obtain cutoffs from outcome_constraints. - if metric_signature not in optimization_config.objective.metric_signatures: + if metric_signature not in optimization_config.objective.metric_names: # Get all outcome constraints for `metric_signature`` that aren't scalarized. return _obtain_cutoffs_from_outcome_constraints( optimization_config=optimization_config, @@ -250,14 +248,11 @@ def _get_cutoffs( objective = assert_is_instance( optimization_config.objective, ScalarizedObjective ) - weight = [ - w for m, w in objective.metric_weights if m.signature == metric_signature - ][0] - # Winsorize from above if the weight is positive and minimize is `True` or the - # weight is negative and minimize is `False`. + weight = [w for m, w in objective.metric_weights if m == metric_signature][0] + # metric_weights are sign-encoded: a negative weight means minimize. return _get_auto_winsorization_cutoffs_single_objective( metric_values=metric_values, - minimize=objective.minimize if weight >= 0 else not objective.minimize, + minimize=weight < 0, ) # Single-objective @@ -302,11 +297,10 @@ def _get_auto_winsorization_cutoffs_multi_objective( AxOptimizationWarning, stacklevel=3, ) - objectives = assert_is_instance(optimization_config.objective, MultiObjective) minimize = [ - objective.minimize - for objective in objectives.objectives - if objective.metric.signature == metric_signature + w < 0 + for name, w in optimization_config.objective.metric_weights + if name == metric_signature ][0] return _get_auto_winsorization_cutoffs_single_objective( metric_values=metric_values, @@ -323,7 +317,7 @@ def _obtain_cutoffs_from_outcome_constraints( # Check for scalarized outcome constraints for the given metric if any( isinstance(oc, ScalarizedOutcomeConstraint) - and metric_signature in [metric.signature for metric in oc.metrics] + and metric_signature in oc.metric_names for oc in optimization_config.outcome_constraints ): warnings.warn( @@ -351,18 +345,18 @@ def _get_non_scalarized_outcome_constraints( oc for oc in optimization_config.outcome_constraints if not isinstance(oc, ScalarizedOutcomeConstraint) - and oc.metric.signature == metric_signature + and oc.metric_names[0] == metric_signature ] def _get_objective_threshold_from_moo_config( optimization_config: MultiObjectiveOptimizationConfig, metric_signature: str -) -> list[ObjectiveThreshold]: +) -> list[OutcomeConstraint]: """Get the non-relative objective threshold for a given metric.""" return [ ot for ot in optimization_config.objective_thresholds - if ot.metric.signature == metric_signature + if ot.metric_names[0] == metric_signature ] diff --git a/ax/analysis/diagnostics.py b/ax/analysis/diagnostics.py index ade5adad082..46feb451ebc 100644 --- a/ax/analysis/diagnostics.py +++ b/ax/analysis/diagnostics.py @@ -70,7 +70,7 @@ def compute( metric_names = list(experiment.metrics.keys()) else: # Extract all metric names from the OptimizationConfig. - metric_names = [*none_throws(experiment.optimization_config).metrics.keys()] + metric_names = [*none_throws(experiment.optimization_config).metric_names] is_bandit = generation_strategy and is_bandit_experiment( generation_strategy_name=generation_strategy.name diff --git a/ax/analysis/healthcheck/constraints_feasibility.py b/ax/analysis/healthcheck/constraints_feasibility.py index 11579c4c3ac..ae4bda5e311 100644 --- a/ax/analysis/healthcheck/constraints_feasibility.py +++ b/ax/analysis/healthcheck/constraints_feasibility.py @@ -117,7 +117,7 @@ def validate_applicable_state( ): optimization_config = none_throws(experiment.optimization_config) constraint_metric_names = [ - outcome_constraint.metric.name + outcome_constraint.metric_names[0] for outcome_constraint in optimization_config.outcome_constraints ] @@ -185,7 +185,7 @@ def compute( arm_data = prepare_arm_data( experiment=experiment, - metric_names=[*optimization_config.metrics.keys()], + metric_names=[*optimization_config.metric_names], use_model_predictions=True, adapter=relevant_adapter, trial_statuses=[TrialStatus.COMPLETED, TrialStatus.RUNNING], @@ -221,7 +221,7 @@ def compute( if isinstance(oc, ScalarizedOutcomeConstraint): constraint_map[str(oc)] = oc else: - constraint_map[oc.metric.name] = oc + constraint_map[oc.metric_names[0]] = oc for col in p_feasible_cols: constraint_name = col.replace("p_feasible_", "") diff --git a/ax/analysis/healthcheck/early_stopping_healthcheck.py b/ax/analysis/healthcheck/early_stopping_healthcheck.py index 49bfca2e30c..9af56cb297f 100644 --- a/ax/analysis/healthcheck/early_stopping_healthcheck.py +++ b/ax/analysis/healthcheck/early_stopping_healthcheck.py @@ -404,7 +404,7 @@ def _report_early_stopping_nudge( default early stopping strategy is available. """ opt_config = none_throws(experiment.optimization_config) - metric = next(iter(opt_config.objective.metrics)) + metric = experiment.get_metric(opt_config.objective.metric_names[0]) try: savings = estimate_hypothetical_early_stopping_savings( experiment=experiment, @@ -492,9 +492,7 @@ def _get_map_metrics(self, experiment: Experiment) -> list[MapMetric]: # Sort so that objective metrics appear first if experiment.optimization_config is not None: - metric_names = [ - m.name for m in experiment.optimization_config.objective.metrics - ] + metric_names = experiment.optimization_config.objective.metric_names map_metrics.sort( key=lambda e: e.name in metric_names, reverse=True, diff --git a/ax/analysis/healthcheck/tests/test_complexity_rating.py b/ax/analysis/healthcheck/tests/test_complexity_rating.py index 61f4ff75cbc..d632d3a625d 100644 --- a/ax/analysis/healthcheck/tests/test_complexity_rating.py +++ b/ax/analysis/healthcheck/tests/test_complexity_rating.py @@ -99,6 +99,9 @@ def test_objectives_count(self) -> None: for num_objs, expected_status, expected_tier, expected_msg in test_cases: with self.subTest(num_objectives=num_objs): metrics = [Metric(name=f"m{i}") for i in range(num_objs)] + for m in metrics: + if m.name not in self.experiment.metrics: + self.experiment.add_tracking_metric(m) self.experiment._optimization_config = MultiObjectiveOptimizationConfig( objective=MultiObjective( objectives=[ @@ -116,7 +119,11 @@ def test_objectives_count(self) -> None: def test_constraints(self) -> None: with self.subTest(constraint_type="outcome"): + obj_metric = Metric(name="obj") metrics = [Metric(name=f"m{i}") for i in range(3)] + for m in [obj_metric] + metrics: + if m.name not in self.experiment.metrics: + self.experiment.add_tracking_metric(m) self.experiment._optimization_config = OptimizationConfig( objective=Objective(metric=Metric(name="obj"), minimize=False), outcome_constraints=[ diff --git a/ax/analysis/insights.py b/ax/analysis/insights.py index 4ae61e17b98..a8a28693132 100644 --- a/ax/analysis/insights.py +++ b/ax/analysis/insights.py @@ -77,7 +77,7 @@ def compute( objective_names = optimization_config.objective.metric_names constraint_names = [ - constraint.metric.name + constraint.metric_names[0] for constraint in optimization_config.outcome_constraints ] # Relativize the effects if the status quo is set and there are BatchTrials diff --git a/ax/analysis/plotly/constraint_feasibility.py b/ax/analysis/plotly/constraint_feasibility.py index c0546edf59d..fb8789259de 100644 --- a/ax/analysis/plotly/constraint_feasibility.py +++ b/ax/analysis/plotly/constraint_feasibility.py @@ -107,7 +107,7 @@ def validate_applicable_state( return outcome_constraints_validation_str outcome_constraint_metrics = [ - outcome_constraint.metric.name + outcome_constraint.metric_names[0] for outcome_constraint in none_throws( experiment.optimization_config ).outcome_constraints @@ -153,12 +153,10 @@ def compute( # Collect all metrics and compute constraint_metric_names = [ - outcome_constraint.metric.name + outcome_constraint.metric_names[0] for outcome_constraint in optimization_config.outcome_constraints ] - objective_metric_names = [ - metric.name for metric in optimization_config.objective.metrics - ] + objective_metric_names = optimization_config.objective.metric_names all_metric_names = list(set(constraint_metric_names + objective_metric_names)) arm_data = prepare_arm_data( @@ -244,7 +242,7 @@ def _get_constraint_names( if isinstance(oc, ScalarizedOutcomeConstraint): constraint_names.append(str(oc)) else: - constraint_names.append(oc.metric.name) + constraint_names.append(oc.metric_names[0]) return constraint_names diff --git a/ax/analysis/plotly/objective_p_feasible_frontier.py b/ax/analysis/plotly/objective_p_feasible_frontier.py index f160eaaae45..bace1a802ac 100644 --- a/ax/analysis/plotly/objective_p_feasible_frontier.py +++ b/ax/analysis/plotly/objective_p_feasible_frontier.py @@ -211,7 +211,7 @@ def compute( df = prepare_arm_data( experiment=experiment, - metric_names=[*optimization_config.metrics.keys()], + metric_names=[*optimization_config.metric_names], adapter=relevant_adapter, use_model_predictions=True, relativize=self.relativize, @@ -220,7 +220,7 @@ def compute( trial_statuses=self.trial_statuses, ) - objective_name = optimization_config.objective.metric.name + objective_name = optimization_config.objective.metric_names[0] fig = _prepare_figure_scatter( df=df, diff --git a/ax/analysis/plotly/p_feasible.py b/ax/analysis/plotly/p_feasible.py index b051ba5e02c..a20325f8f15 100644 --- a/ax/analysis/plotly/p_feasible.py +++ b/ax/analysis/plotly/p_feasible.py @@ -86,7 +86,7 @@ def validate_applicable_state( if optimization_config is None: return "Experiment must have an OptimizationConfig." - if len(optimization_config.objective.metrics) > 1: + if optimization_config.objective.is_multi_objective: return "Only single-objective optimization is supported." outcome_constraints_validation_str = validate_outcome_constraints( @@ -96,7 +96,7 @@ def validate_applicable_state( return outcome_constraints_validation_str outcome_constraint_metrics = [ - outcome_constraint.metric.name + outcome_constraint.metric_names[0] for outcome_constraint in optimization_config.outcome_constraints ] @@ -151,7 +151,7 @@ def compute( arm_data = prepare_arm_data( experiment=experiment, metric_names=[ - outcome_constrinat.metric.name + outcome_constrinat.metric_names[0] for outcome_constrinat in optimization_config.outcome_constraints ], use_model_predictions=self.use_model_predictions, diff --git a/ax/analysis/plotly/tests/test_marginal_effects.py b/ax/analysis/plotly/tests/test_marginal_effects.py index 107381844cd..ab9b2cf49dc 100644 --- a/ax/analysis/plotly/tests/test_marginal_effects.py +++ b/ax/analysis/plotly/tests/test_marginal_effects.py @@ -30,7 +30,7 @@ def setUp(self) -> None: name="test_experiment", is_test=True, optimization_config=OptimizationConfig( - objective=Objective(Metric("metric_1", lower_is_better=False)) + objective=Objective(metric=Metric("metric_1", lower_is_better=False)) ), ) num_arms = 3 diff --git a/ax/analysis/plotly/tests/test_objective_p_feasible_frontier.py b/ax/analysis/plotly/tests/test_objective_p_feasible_frontier.py index 7405279658d..37fac1dbbd6 100644 --- a/ax/analysis/plotly/tests/test_objective_p_feasible_frontier.py +++ b/ax/analysis/plotly/tests/test_objective_p_feasible_frontier.py @@ -18,7 +18,7 @@ MultiObjectiveOptimizationConfig, OptimizationConfig, ) -from ax.core.outcome_constraint import ScalarizedOutcomeConstraint +from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus from ax.core.types import ComparisonOp from ax.utils.common.testutils import TestCase @@ -43,9 +43,11 @@ def setUp(self) -> None: objective=Objective(metric=self.experiment.metrics["branin_a"]), outcome_constraints=self.experiment.optimization_config.outcome_constraints, ) - oc = none_throws(self.experiment.optimization_config).outcome_constraints[0] - oc.bound = 10.0 - oc.op = ComparisonOp.LEQ + opt_config = none_throws(self.experiment.optimization_config) + oc = opt_config.outcome_constraints[0] + opt_config.outcome_constraints[0] = OutcomeConstraint( + expression=f"{oc.metric_names[0]} <= 10.0", + ) def test_trial_statuses_behavior(self) -> None: # When neither trial_statuses nor trial_index is provided, diff --git a/ax/analysis/plotly/utility_progression.py b/ax/analysis/plotly/utility_progression.py index d1317a617fc..5ffc765d761 100644 --- a/ax/analysis/plotly/utility_progression.py +++ b/ax/analysis/plotly/utility_progression.py @@ -19,7 +19,6 @@ ) from ax.analysis.utils import validate_experiment from ax.core.experiment import Experiment -from ax.core.objective import ScalarizedObjective from ax.exceptions.core import ExperimentNotReadyError from ax.generation_strategy.generation_strategy import GenerationStrategy from ax.service.utils.best_point import get_trace @@ -159,23 +158,22 @@ def compute( ) else: objective = none_throws(opt_config).objective - minimize = objective.minimize - direction = "minimize" if minimize else "maximize" # Handle ScalarizedObjective vs regular Objective - if isinstance(objective, ScalarizedObjective): + if objective.is_scalarized_objective: expression = objective.expression y_label = f"Best Observed {expression}" subtitle = ( f"Shows the best scalarized objective value (formula: " - f"{expression}) achieved so far across completed trials " - f"(objective is to {direction}). {_TRACE_INDEX_EXPLANATION} " + f"{expression}) achieved so far across completed trials. " + f"{_TRACE_INDEX_EXPLANATION} " f"{_CUMULATIVE_BEST_EXPLANATION} " f"{_INFEASIBLE_TRIALS_EXPLANATION}" ) else: - # Regular single-objective - objective_name = objective.metric.name + minimize = objective.minimize + direction = "minimize" if minimize else "maximize" + objective_name = objective.metric_names[0] y_label = f"Best Observed {objective_name}" subtitle = ( f"Shows the best {objective_name} value achieved so far across " diff --git a/ax/analysis/plotly/utils.py b/ax/analysis/plotly/utils.py index 986f6cd523f..0114d95489f 100644 --- a/ax/analysis/plotly/utils.py +++ b/ax/analysis/plotly/utils.py @@ -12,7 +12,7 @@ import pandas as pd from ax.analysis.plotly.color_constants import BOTORCH_COLOR_SCALE, LIGHT_AX_BLUE from ax.core.experiment import Experiment -from ax.core.objective import MultiObjective, ScalarizedObjective +from ax.core.objective import ScalarizedObjective from ax.core.trial_status import DEFAULT_ANALYSIS_STATUSES, TrialStatus from ax.exceptions.core import UnsupportedError from plotly import express as px @@ -208,7 +208,7 @@ def select_metric(experiment: Experiment) -> str: "Cannot infer metric to plot from Experiment without OptimizationConfig" ) objective = experiment.optimization_config.objective - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: raise UnsupportedError( "Cannot infer metric to plot from MultiObjective, please specify a metric" ) @@ -217,7 +217,7 @@ def select_metric(experiment: Experiment) -> str: "Cannot infer metric to plot from ScalarizedObjective, please " "specify a metric" ) - return experiment.optimization_config.objective.metric.name + return experiment.optimization_config.objective.metric_names[0] def get_trial_statuses_with_fallback( diff --git a/ax/analysis/results.py b/ax/analysis/results.py index deda3d0f4af..a5481e8e133 100644 --- a/ax/analysis/results.py +++ b/ax/analysis/results.py @@ -95,10 +95,7 @@ def compute( if (optimization_config := experiment.optimization_config) is not None: objective_names = optimization_config.objective.metric_names for oc in optimization_config.outcome_constraints: - if isinstance(oc, ScalarizedOutcomeConstraint): - constraint_names.extend([m.name for m in oc.metrics]) - else: - constraint_names.append(oc.metric.name) + constraint_names.extend(oc.metric_names) relevant_adapter = extract_relevant_adapter( experiment=experiment, diff --git a/ax/analysis/tests/test_metric_summary.py b/ax/analysis/tests/test_metric_summary.py index 9cf43b292cf..0c16ccd1bd8 100644 --- a/ax/analysis/tests/test_metric_summary.py +++ b/ax/analysis/tests/test_metric_summary.py @@ -27,7 +27,7 @@ def test_compute(self) -> None: # TODO: Debug error raised by # client.configure_metrics(metrics=[IMetric(name="qux")]) - client._experiment._tracking_metrics = {"qux": Metric(name="qux")} + client._experiment._metrics["qux"] = Metric(name="qux") analysis = MetricSummary() diff --git a/ax/analysis/tests/test_results.py b/ax/analysis/tests/test_results.py index 411a62bdab9..03662837f2b 100644 --- a/ax/analysis/tests/test_results.py +++ b/ax/analysis/tests/test_results.py @@ -329,12 +329,9 @@ def test_compute_with_status_quo_relativizes(self) -> None: def test_compute_without_optimization_config(self) -> None: # Setup: Create experiment without optimization config experiment = get_branin_experiment() - metrics = [ - m.clone() - for m in none_throws(experiment.optimization_config).metrics.values() - ] experiment._optimization_config = None - experiment.add_tracking_metrics(metrics) + # Metrics are already in experiment._metrics; with no opt config, + # they are automatically treated as tracking metrics. trial = experiment.new_trial() trial.add_arm( diff --git a/ax/analysis/trials.py b/ax/analysis/trials.py index bf6c3753b1b..f1d38aa2e4b 100644 --- a/ax/analysis/trials.py +++ b/ax/analysis/trials.py @@ -133,7 +133,7 @@ def compute( if (optimization_config := experiment.optimization_config) is not None: objective_names = optimization_config.objective.metric_names constraint_names = [ - constraint.metric.name + constraint.metric_names[0] for constraint in optimization_config.outcome_constraints ] diff --git a/ax/analysis/utils.py b/ax/analysis/utils.py index 4450f3f6616..1e6e50eab22 100644 --- a/ax/analysis/utils.py +++ b/ax/analysis/utils.py @@ -598,9 +598,9 @@ def _get_scalarized_constraint_mean_and_sem( combined_var = np.zeros(n_rows) all_metrics_present = True - for metric, weight in constraint.metric_weights: - mean_col = f"{metric.name}_mean" - sem_col = f"{metric.name}_sem" + for metric_name, weight in constraint.metric_weights: + mean_col = f"{metric_name}_mean" + sem_col = f"{metric_name}_sem" if mean_col in df.columns: combined_mean += weight * df[mean_col].values @@ -663,7 +663,7 @@ def _prepare_p_feasible( means.append(mean.tolist()) sigmas.append(sem.tolist()) else: - metric_name = oc.metric.name + metric_name = oc.metric_names[0] if f"{metric_name}_mean" in df_constraint.columns: means.append(df_constraint[f"{metric_name}_mean"].tolist()) else: @@ -749,7 +749,7 @@ def _prepare_p_feasible_per_constraint( mean, sigma = _get_scalarized_constraint_mean_and_sem(df_constraint, oc) oc_display_name = str(oc) else: - metric_name = oc.metric.name + metric_name = oc.metric_names[0] oc_display_name = metric_name if f"{metric_name}_mean" in df_constraint.columns: @@ -1139,7 +1139,7 @@ def validate_outcome_constraints( return "Experiment must have an OptimizationConfig." outcome_constraint_metrics = [ - outcome_constraint.metric.name + outcome_constraint.metric_names[0] for outcome_constraint in optimization_config.outcome_constraints ] if len(outcome_constraint_metrics) == 0: diff --git a/ax/api/client.py b/ax/api/client.py index 8343a4e6c33..ee2d98882bb 100644 --- a/ax/api/client.py +++ b/ax/api/client.py @@ -32,8 +32,9 @@ from ax.core.analysis_card import AnalysisCardBase from ax.core.arm import Arm from ax.core.experiment import Experiment +from ax.core.map_metric import MapMetric from ax.core.metric import Metric -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.observation import ObservationFeatures from ax.core.optimization_config import OptimizationConfig from ax.core.runner import Runner @@ -200,6 +201,37 @@ def configure_optimization( if pruning_target_arm is not None: optimization_config.pruning_target_parameterization = pruning_target_arm + # Register any new metrics from the optimization config on the experiment + # before setting the optimization config, which validates that all + # referenced metrics exist on the experiment. + # Determine lower_is_better for objective metrics from weights. + obj = optimization_config.objective + objective_lower_is_better: dict[str, bool] = {} + if obj.is_multi_objective: + parts = [p.strip() for p in obj.expression.split(",")] + for part in parts: + sub = Objective(expression=part) + for name, weight in sub.metric_weights: + objective_lower_is_better[name] = weight < 0 + else: + for name, weight in obj.metric_weights: + objective_lower_is_better[name] = weight < 0 + + # Register objective metrics first (preserving expression order), + # then constraint metrics. + all_metric_names: list[str] = list(obj.metric_names) + for oc in optimization_config.all_constraints: + for name in oc.metric_names: + if name not in all_metric_names: + all_metric_names.append(name) + + for metric_name in all_metric_names: + if metric_name not in self._experiment.metrics: + lower_is_better = objective_lower_is_better.get(metric_name, None) + self._experiment.add_tracking_metric( + MapMetric(name=metric_name, lower_is_better=lower_is_better) + ) + self._experiment.optimization_config = optimization_config self._set_metrics(metrics=list(old_metrics.values())) self._save_experiment_to_db_if_possible(experiment=self._experiment) @@ -513,7 +545,7 @@ def complete_trial( # Log metric availability for user visibility. if (optimization_config := self._experiment.optimization_config) is not None: trial_data = self._experiment.lookup_data(trial_indices=[trial_index]) - missing_metrics = set(optimization_config.metrics.keys()) - { + missing_metrics = optimization_config.metric_names - { *trial_data.metric_names } @@ -1220,52 +1252,14 @@ def _overwrite_metric(self, metric: Metric) -> None: handled in self._set_metrics). """ - # Check the OptimizationConfig first - if (optimization_config := self._experiment.optimization_config) is not None: - # Check the objective - if isinstance( - multi_objective := optimization_config.objective, MultiObjective - ): - for i in range(len(multi_objective.objectives)): - if metric.name == multi_objective.objectives[i].metric.name: - multi_objective._objectives[i]._metric = metric - return - elif isinstance( - scalarized_objective := optimization_config.objective, - ScalarizedObjective, - ): - for i in range(len(scalarized_objective.metrics)): - if metric.name == scalarized_objective.metrics[i].name: - scalarized_objective._metrics[i] = metric - return - elif ( - isinstance(optimization_config.objective, Objective) - and metric.name == optimization_config.objective.metric.name - ): - optimization_config.objective._metric = metric - return - - # Check the outcome constraints - for i in range(len(optimization_config.outcome_constraints)): - if ( - metric.name - == optimization_config.outcome_constraints[i].metric.name - ): - optimization_config._outcome_constraints[i]._metric = metric - return - - # Check the tracking metrics - if metric.name in self._experiment._tracking_metrics.keys(): - self._experiment._tracking_metrics[metric.name] = metric - return - - # If an equivalently named Metric does not exist, add it as a tracking - # metric. - self._experiment.add_tracking_metric(metric=metric) - logger.warning( - f"Metric {metric} not found in optimization config, added as tracking " - "metric." - ) + if metric.name in self._experiment.metrics: + self._experiment.update_metric(metric) + else: + self._experiment.add_tracking_metric(metric=metric) + logger.warning( + f"Metric {metric} not found on experiment, added as a new " + "tracking metric." + ) # -------------------- Section 5.3: Storage utilies ------------------------------- def _to_json_snapshot(self) -> dict[str, Any]: diff --git a/ax/api/tests/test_client.py b/ax/api/tests/test_client.py index afa38da272e..7a62d233f02 100644 --- a/ax/api/tests/test_client.py +++ b/ax/api/tests/test_client.py @@ -26,7 +26,7 @@ from ax.core.data import Data from ax.core.experiment import Experiment from ax.core.map_metric import MapMetric -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.optimization_config import OptimizationConfig from ax.core.outcome_constraint import ComparisonOp, OutcomeConstraint from ax.core.parameter import ( @@ -223,10 +223,12 @@ def test_configure_optimization(self) -> None: optimization_config.objective, Objective, ) - self.assertEqual(objective.metric.name, "metric1") + self.assertEqual(objective.metric_names[0], "metric1") # Verify no metrics were removed, just moved from tracking to objective - all_metrics = [objective.metric] + list(client._experiment.tracking_metrics) + all_metrics = [client._experiment.get_metric(objective.metric_names[0])] + list( + client._experiment.tracking_metrics + ) self.assertEqual(len(all_metrics), 3) def test_configure_optimization_with_pruning_target_parameterization(self) -> None: @@ -367,30 +369,32 @@ def test_configure_metric(self) -> None: self.assertEqual( custom_metric, - none_throws(client._experiment.optimization_config).objective.metric, + client._experiment.get_metric("custom"), ) # Test replacing a multi-objective client.configure_optimization(objective="custom, foo") client.configure_metrics(metrics=[custom_metric]) - self.assertIn( + self.assertEqual( custom_metric, - assert_is_instance( - none_throws(client._experiment.optimization_config).objective, - MultiObjective, - ).metrics, + client._experiment.get_metric("custom"), + ) + self.assertIn( + "custom", + none_throws(client._experiment.optimization_config).objective.metric_names, ) # Test replacing a scalarized objective client.configure_optimization(objective="custom + foo") client.configure_metrics(metrics=[custom_metric]) - self.assertIn( + self.assertEqual( custom_metric, - assert_is_instance( - none_throws(client._experiment.optimization_config).objective, - ScalarizedObjective, - ).metrics, + client._experiment.get_metric("custom"), + ) + self.assertIn( + "custom", + none_throws(client._experiment.optimization_config).objective.metric_names, ) # Test replacing an outcome constraint @@ -401,9 +405,7 @@ def test_configure_metric(self) -> None: self.assertEqual( custom_metric, - none_throws(client._experiment.optimization_config) - .outcome_constraints[0] - .metric, + client._experiment.get_metric("custom"), ) # Test adding a tracking metric @@ -1622,32 +1624,36 @@ def test_overwrite_metric(self) -> None: client.configure_metrics(metrics=[foo_metric]) objective = none_throws(client._experiment.optimization_config).objective self.assertIsInstance(objective, Objective) - self.assertIs(objective.metric, foo_metric) + self.assertIs( + client._experiment.get_metric(objective.metric_names[0]), foo_metric + ) # test 2: Replace the outcome-constraint metric bar_metric = DummyMetric(name="bar") client.configure_metrics(metrics=[bar_metric]) oc = none_throws(client._experiment.optimization_config).outcome_constraints[0] - self.assertIs(oc.metric, bar_metric) + self.assertIs(client._experiment.get_metric(oc.metric_names[0]), bar_metric) # test 3: Add a tracking metric, then replace it by name baz_metric_1 = DummyMetric(name="baz") client.configure_metrics(metrics=[baz_metric_1]) - self.assertIn("baz", client._experiment._tracking_metrics) - self.assertIs(client._experiment._tracking_metrics["baz"], baz_metric_1) + tracking_names = {m.name for m in client._experiment.tracking_metrics} + self.assertIn("baz", tracking_names) + self.assertIs(client._experiment.get_metric("baz"), baz_metric_1) baz_metric_2 = DummyMetric(name="baz") client.configure_metrics(metrics=[baz_metric_2]) - self.assertIs(client._experiment._tracking_metrics["baz"], baz_metric_2) + self.assertIs(client._experiment.get_metric("baz"), baz_metric_2) # test 4: Metric name not present anywhere, should be added as tracking + warn quux_metric = DummyMetric(name="quux") with self.assertLogs(logger="ax.api.client", level="WARNING") as logs: client.configure_metrics(metrics=[quux_metric]) - self.assertIn("quux", client._experiment._tracking_metrics) - self.assertIs(client._experiment._tracking_metrics["quux"], quux_metric) + tracking_names = {m.name for m in client._experiment.tracking_metrics} + self.assertIn("quux", tracking_names) + self.assertIs(client._experiment.get_metric("quux"), quux_metric) self.assertTrue( - any("added as tracking metric" in msg for msg in logs.output), + any("added as a new tracking metric" in msg for msg in logs.output), "Expected a warning that the metric was added as a tracking metric.", ) @@ -1655,21 +1661,13 @@ def test_overwrite_metric(self) -> None: client.configure_optimization(objective="foo, qux") qux_metric_moo = DummyMetric(name="qux") client.configure_metrics(metrics=[qux_metric_moo]) - moo = assert_is_instance( - none_throws(client._experiment.optimization_config).objective, - MultiObjective, - ) - self.assertIn(qux_metric_moo, moo.metrics) + self.assertEqual(qux_metric_moo, client._experiment.get_metric("qux")) # test 6: Replace inside a ScalarizedObjective client.configure_optimization(objective="foo + qux") qux_metric_scalar = DummyMetric(name="qux") client.configure_metrics(metrics=[qux_metric_scalar]) - scalar = assert_is_instance( - none_throws(client._experiment.optimization_config).objective, - ScalarizedObjective, - ) - self.assertIn(qux_metric_scalar, scalar.metrics) + self.assertEqual(qux_metric_scalar, client._experiment.get_metric("qux")) def test_configure_generation_strategy_with_simplify(self) -> None: client = Client() diff --git a/ax/api/utils/instantiation/from_string.py b/ax/api/utils/instantiation/from_string.py index 3e72a097e08..989622208e8 100644 --- a/ax/api/utils/instantiation/from_string.py +++ b/ax/api/utils/instantiation/from_string.py @@ -7,27 +7,13 @@ from collections.abc import Sequence -from ax.core.map_metric import MapMetric -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, OptimizationConfig, ) -from ax.core.outcome_constraint import ( - ComparisonOp, - ObjectiveThreshold, - OutcomeConstraint, - ScalarizedOutcomeConstraint, -) +from ax.core.outcome_constraint import OutcomeConstraint from ax.exceptions.core import UserInputError -from ax.utils.common.string_utils import sanitize_name, unsanitize_name -from ax.utils.common.sympy import extract_coefficient_dict_from_inequality -from pyre_extensions import assert_is_instance, none_throws -from sympy.core.add import Add -from sympy.core.expr import Expr -from sympy.core.mul import Mul -from sympy.core.symbol import Symbol -from sympy.core.sympify import sympify def optimization_config_from_string( @@ -40,34 +26,25 @@ def optimization_config_from_string( the multi-objective case where they will be converted to objective thresholds. """ - objective = parse_objective(objective_str=objective_str) + objective = Objective(expression=objective_str) - if outcome_constraint_strs is not None: - outcome_constraints = [ - parse_outcome_constraint(constraint_str=constraint_str) - for constraint_str in outcome_constraint_strs - ] - else: - outcome_constraints = None + outcome_constraints: list[OutcomeConstraint] | None = ( + [OutcomeConstraint(expression=s) for s in outcome_constraint_strs] + if outcome_constraint_strs is not None + else None + ) - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: # Convert OutcomeConstraints to ObjectiveThresholds if relevant - objective_metric_names = {metric.name for metric in objective.metrics} + objective_metric_names = set(objective.metric_names) true_outcome_constraints = [] - objective_thresholds: list[ObjectiveThreshold] = [] + objective_thresholds: list[OutcomeConstraint] = [] for outcome_constraint in outcome_constraints or []: if ( - not isinstance(outcome_constraint, ScalarizedOutcomeConstraint) - and outcome_constraint.metric.name in objective_metric_names + len(outcome_constraint.metric_names) == 1 + and outcome_constraint.metric_names[0] in objective_metric_names ): - objective_thresholds.append( - ObjectiveThreshold( - metric=outcome_constraint.metric, - bound=outcome_constraint.bound, - relative=outcome_constraint.relative, - op=outcome_constraint.op, - ) - ) + objective_thresholds.append(outcome_constraint) else: true_outcome_constraints.append(outcome_constraint) @@ -78,9 +55,9 @@ def optimization_config_from_string( ) # Ensure that outcome constraints are not placed on the objective metric - objective_metric_names = {metric.name for metric in objective.metrics} + objective_metric_names = set(objective.metric_names) for outcome_constraint in outcome_constraints or []: - if outcome_constraint.metric.name in objective_metric_names: + if outcome_constraint.metric_names[0] in objective_metric_names: raise UserInputError( "Outcome constraints may not be placed on the objective metric " f"except in the multi-objective case, found {objective_str} and " @@ -91,128 +68,3 @@ def optimization_config_from_string( objective=objective, outcome_constraints=outcome_constraints, ) - - -def parse_objective(objective_str: str) -> Objective: - """ - Parse an objective string into an Objective object using SymPy. - - Currently only supports linear objectives of the form "a * x + b * y" and tuples of - linear objectives. - """ - # Parse the objective string into a SymPy expression - expression = sympify(sanitize_name(objective_str)) - - if isinstance(expression, tuple): # Multi-objective - return MultiObjective( - objectives=[ - _create_single_objective(expression=term) for term in expression - ] - ) - - return _create_single_objective(expression=assert_is_instance(expression, Expr)) - - -def parse_outcome_constraint(constraint_str: str) -> OutcomeConstraint: - """ - Parse an outcome constraint string into an OutcomeConstraint object using SymPy. - Currently only supports linear constraints of the form "a * x + b * y >= k" or - "a * x + b * y <= k". - - To indicate a relative constraint (i.e. performance relative to some baseline) - multiply your bound by "baseline". For example "qps >= 0.95 * baseline" will - constrain such that the QPS is at least 95% of the baseline arm's QPS. - """ - coefficient_dict = extract_coefficient_dict_from_inequality( - inequality_str=constraint_str - ) - - # Iterate through the coefficients to extract the parameter names and weights and - # the bound - constraint_dict: dict[str, float] = {} - bound = 0 - is_relative = False - for term, coefficient in coefficient_dict.items(): - if term.is_symbol: - if term.name == "baseline": - # Invert because we are "moving" the bound to the right hand side - bound = -1 * coefficient - is_relative = True - else: - constraint_dict[term.name] = coefficient - elif term.is_number: - # Invert because we are "moving" the bound to the right hand side - bound = -1 * coefficient - else: - raise UserInputError( - f"Only linear outcome constraints are supported, found {constraint_str}" - ) - - if len(constraint_dict) == 1: - term, coefficient = next(iter(constraint_dict.items())) - - return OutcomeConstraint( - metric=MapMetric(name=unsanitize_name(term)), - op=ComparisonOp.LEQ if coefficient > 0 else ComparisonOp.GEQ, - bound=bound / coefficient, - relative=is_relative, - ) - - names, coefficients = zip(*constraint_dict.items()) - return ScalarizedOutcomeConstraint( - metrics=[MapMetric(name=unsanitize_name(name)) for name in names], - op=ComparisonOp.LEQ, - weights=[*coefficients], - bound=bound, - relative=is_relative, - ) - - -def _create_single_objective(expression: Expr) -> Objective: - """ - Create an Objective or ScalarizedObjective from a linear SymPy expression. - - All expressions are assumed to represent maximization objectives. - """ - - # If the expression is a just a Symbol it represents a single metric objective - if isinstance(expression, Symbol): - return Objective( - metric=MapMetric( - name=unsanitize_name(str(expression.name)), lower_is_better=False - ), - minimize=False, - ) - - # If the expression is a Mul it likely represents a single metric objective but - # some additional validation is required - if isinstance(expression, Mul): - symbol, *other_symbols = expression.free_symbols - if len(other_symbols) > 0: - raise UserInputError( - f"Only linear objectives are supported, found {expression}." - ) - - # Since the objectives 1 * loss and 2 * loss are equivalent, we can just use - # the sign from the coefficient rather than its value - minimize = bool( - none_throws(expression.as_coefficient(assert_is_instance(symbol, Expr))) < 0 - ) - - return Objective( - metric=MapMetric( - name=unsanitize_name(str(symbol)), lower_is_better=minimize - ), - minimize=minimize, - ) - - # If the expression is an Add it represents a scalarized objective - elif isinstance(expression, Add): - names, coefficients = zip(*expression.as_coefficients_dict().items()) - return ScalarizedObjective( - metrics=[MapMetric(name=unsanitize_name(str(name))) for name in names], - weights=[float(coefficient) for coefficient in coefficients], - minimize=False, - ) - - raise UserInputError(f"Only linear objectives are supported, found {expression}.") diff --git a/ax/api/utils/instantiation/tests/test_from_string.py b/ax/api/utils/instantiation/tests/test_from_string.py index 3dfff31f277..1a116e5ccff 100644 --- a/ax/api/utils/instantiation/tests/test_from_string.py +++ b/ax/api/utils/instantiation/tests/test_from_string.py @@ -4,23 +4,13 @@ # LICENSE file in the root directory of this source tree. # pyre-strict -from ax.api.utils.instantiation.from_string import ( - optimization_config_from_string, - parse_objective, - parse_outcome_constraint, -) -from ax.core.map_metric import MapMetric -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.api.utils.instantiation.from_string import optimization_config_from_string +from ax.core.objective import Objective from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, OptimizationConfig, ) -from ax.core.outcome_constraint import ( - ComparisonOp, - ObjectiveThreshold, - OutcomeConstraint, - ScalarizedOutcomeConstraint, -) +from ax.core.outcome_constraint import OutcomeConstraint from ax.exceptions.core import UserInputError from ax.utils.common.testutils import TestCase @@ -31,9 +21,7 @@ def test_optimization_config_from_string(self) -> None: self.assertEqual( only_objective, OptimizationConfig( - objective=Objective( - metric=MapMetric(name="ne", lower_is_better=False), minimize=False - ), + objective=Objective(expression="ne"), ), ) @@ -43,16 +31,9 @@ def test_optimization_config_from_string(self) -> None: self.assertEqual( with_constraints, OptimizationConfig( - objective=Objective( - metric=MapMetric(name="ne", lower_is_better=False), minimize=False - ), + objective=Objective(expression="ne"), outcome_constraints=[ - OutcomeConstraint( - metric=MapMetric(name="qps"), - op=ComparisonOp.GEQ, - bound=0.0, - relative=False, - ) + OutcomeConstraint(expression="qps >= 0"), ], ), ) @@ -64,140 +45,20 @@ def test_optimization_config_from_string(self) -> None: self.assertEqual( with_constraints_and_objective_threshold, MultiObjectiveOptimizationConfig( - objective=MultiObjective( - objectives=[ - Objective( - metric=MapMetric(name="ne", lower_is_better=True), - minimize=True, - ), - Objective( - metric=MapMetric(name="qps", lower_is_better=False), - minimize=False, - ), - ] - ), + objective=Objective(expression="-ne, qps"), outcome_constraints=[ - OutcomeConstraint( - metric=MapMetric(name="flops"), - op=ComparisonOp.LEQ, - bound=1000000.0, - relative=False, - ) + OutcomeConstraint(expression="flops <= 1000000"), ], objective_thresholds=[ - ObjectiveThreshold( - metric=MapMetric(name="qps"), - op=ComparisonOp.GEQ, - bound=1000.0, - relative=False, - ) - ], - ), - ) - - def test_parse_objective(self) -> None: - single_objective = parse_objective(objective_str="ne") - self.assertEqual( - single_objective, - Objective( - metric=MapMetric(name="ne", lower_is_better=False), minimize=False - ), - ) - - maximize_single_objective = parse_objective(objective_str="-qps") - self.assertEqual( - maximize_single_objective, - Objective( - metric=MapMetric(name="qps", lower_is_better=True), minimize=True - ), - ) - - scalarized_objective = parse_objective( - objective_str="0.5 * ne1 + 0.3 * ne2 + 0.2 * ne3" - ) - self.assertEqual( - scalarized_objective, - ScalarizedObjective( - metrics=[ - MapMetric(name="ne1"), - MapMetric(name="ne2"), - MapMetric(name="ne3"), + OutcomeConstraint(expression="qps >= 1000"), ], - weights=[0.5, 0.3, 0.2], - minimize=False, - ), - ) - - multiobjective = parse_objective(objective_str="ne, -qps") - self.assertEqual( - multiobjective, - MultiObjective( - objectives=[ - Objective( - metric=MapMetric(name="ne", lower_is_better=False), - minimize=False, - ), - Objective( - metric=MapMetric(name="qps", lower_is_better=True), - minimize=True, - ), - ] ), ) - with self.assertRaisesRegex(UserInputError, "Only linear"): - parse_objective(objective_str="ne * qps") - - def test_parse_outcome_constraint(self) -> None: - constraint = parse_outcome_constraint(constraint_str="flops <= 1000000") - self.assertEqual( - constraint, - OutcomeConstraint( - metric=MapMetric(name="flops"), - op=ComparisonOp.LEQ, - bound=1000000.0, - relative=False, - ), - ) - - flipped_sign = parse_outcome_constraint(constraint_str="flops >= 1000000.0") - self.assertEqual( - flipped_sign, - OutcomeConstraint( - metric=MapMetric(name="flops"), - op=ComparisonOp.GEQ, - bound=1000000.0, - relative=False, - ), - ) - - relative = parse_outcome_constraint(constraint_str="flops <= 105 * baseline") - self.assertEqual( - relative, - OutcomeConstraint( - metric=MapMetric(name="flops"), - op=ComparisonOp.LEQ, - bound=105.0, - relative=True, - ), - ) - - scalarized = parse_outcome_constraint( - constraint_str="0.5 * flops1 + 0.3 * flops2 <= 1000000" - ) - self.assertEqual( - scalarized, - ScalarizedOutcomeConstraint( - metrics=[MapMetric(name="flops1"), MapMetric(name="flops2")], - weights=[0.5, 0.3], - op=ComparisonOp.LEQ, - bound=1000000.0, - relative=False, - ), - ) - - with self.assertRaisesRegex(UserInputError, "Expected an inequality"): - parse_outcome_constraint(constraint_str="flops == 1000000") - - with self.assertRaisesRegex(UserInputError, "Only linear"): - parse_outcome_constraint(constraint_str="flops * flops <= 1000000") + def test_objective_constraint_on_single_objective_raises(self) -> None: + with self.assertRaisesRegex( + UserInputError, "Outcome constraints may not be placed" + ): + optimization_config_from_string( + objective_str="ne", outcome_constraint_strs=["ne >= 0"] + ) diff --git a/ax/benchmark/benchmark.py b/ax/benchmark/benchmark.py index b8b081f8f41..2c5a6496634 100644 --- a/ax/benchmark/benchmark.py +++ b/ax/benchmark/benchmark.py @@ -33,7 +33,7 @@ import numpy.typing as npt import pandas as pd from ax.benchmark.benchmark_method import BenchmarkMethod -from ax.benchmark.benchmark_problem import BenchmarkProblem +from ax.benchmark.benchmark_problem import _is_minimizing, BenchmarkProblem from ax.benchmark.benchmark_result import AggregatedBenchmarkResult, BenchmarkResult from ax.benchmark.benchmark_runner import BenchmarkRunner from ax.benchmark.benchmark_test_function import BenchmarkTestFunction @@ -41,7 +41,7 @@ from ax.core.arm import Arm from ax.core.data import MAP_KEY from ax.core.experiment import Experiment -from ax.core.objective import MultiObjective +from ax.core.metric import Metric from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, OptimizationConfig, @@ -199,6 +199,7 @@ def get_oracle_experiment_from_params( experiment = Experiment( search_space=problem.search_space, optimization_config=problem.optimization_config, + tracking_metrics=(problem.opt_config_metrics or []), ) # The test function produces ground-truth values; noise is handled by @@ -644,7 +645,8 @@ def run_optimization_with_orchestrator( optimization_config=problem.optimization_config, runner=runner, status_quo=sq_arm, - tracking_metrics=problem.tracking_metrics, + tracking_metrics=(problem.opt_config_metrics or []) + + (problem.tracking_metrics or []), auxiliary_experiments_by_purpose=problem.auxiliary_experiments_by_purpose, ) @@ -740,6 +742,7 @@ def compute_baseline_value_from_sobol( test_function: BenchmarkTestFunction, target_fidelity_and_task: Mapping[str, TParamValue] | None = None, n_repeats: int = 50, + opt_config_metrics: list[Metric] | None = None, ) -> float: """ Compute the `baseline_value` that will be assigned to @@ -766,8 +769,8 @@ def compute_baseline_value_from_sobol( # set up a dummy problem so we can use `benchmark_replication` # MOO problems are always higher-is-better because they use hypervolume - higher_is_better = isinstance(optimization_config.objective, MultiObjective) or ( - not optimization_config.objective.minimize + higher_is_better = optimization_config.objective.is_multi_objective or ( + not _is_minimizing(optimization_config.objective) ) dummy_problem = BenchmarkProblem( name="dummy", @@ -781,6 +784,7 @@ def compute_baseline_value_from_sobol( baseline_value=0.0, search_space=search_space, target_fidelity_and_task=target_fidelity_and_task, + opt_config_metrics=opt_config_metrics or [], ) values = np.full(n_repeats, np.nan) diff --git a/ax/benchmark/benchmark_problem.py b/ax/benchmark/benchmark_problem.py index 32ceeefa099..1b571ffaa8f 100644 --- a/ax/benchmark/benchmark_problem.py +++ b/ax/benchmark/benchmark_problem.py @@ -15,20 +15,32 @@ from ax.benchmark.noise import GaussianNoise, Noise from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose from ax.core.metric import Metric -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, - ObjectiveThreshold, OptimizationConfig, ) from ax.core.outcome_constraint import OutcomeConstraint from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter from ax.core.search_space import SearchSpace -from ax.core.types import ComparisonOp, TParamValue +from ax.core.types import TParamValue from ax.exceptions.core import UserInputError from ax.utils.common.base import Base +def _is_minimizing(objective: Objective) -> bool: + """Determine if an objective is minimizing. + + Handles scalarized objectives by checking if all weights are negative + (which indicates minimization in the expression-based API). + """ + if objective.is_scalarized_objective: + return all(w < 0 for _, w in objective.metric_weights) + if objective.is_multi_objective: + return all(w < 0 for _, w in objective.metric_weights) + return objective.minimize + + @dataclass(kw_only=True, repr=True) class BenchmarkProblem(Base): """ @@ -107,6 +119,7 @@ class BenchmarkProblem(Base): dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] | None ) = None tracking_metrics: list[Metric] | None = None + opt_config_metrics: list[Metric] = field(default_factory=list) def __post_init__(self) -> None: # Handle backward compatibility for noise_std parameter @@ -142,7 +155,7 @@ def __post_init__(self) -> None: "The baseline value must be strictly less than the optimal " "value for MOO problems. (These represent hypervolumes.)" ) - elif self.optimization_config.objective.minimize: + elif _is_minimizing(self.optimization_config.objective): if self.baseline_value <= self.optimal_value: raise ValueError( "The baseline value must be strictly greater than the optimal " @@ -166,27 +179,25 @@ def __post_init__(self) -> None: "The worst feasible value must be provided for constrained " "problems (got `None`)" ) - elif self.optimization_config.objective.minimize: - if self.optimal_value > self.worst_feasible_value: + else: + wfv = self.worst_feasible_value + if _is_minimizing(self.optimization_config.objective): + if self.optimal_value > wfv: + raise ValueError( + "The worst feasible value must be greater than or " + "equal to the optimal value for minimization " + "problems." + ) + elif self.optimal_value < wfv: raise ValueError( - "The worst feasible value must be greater than or equal to " - "the optimal value for minimization problems." + "The worst feasible value must be less than or equal to " + "the optimal value for maximization problems." ) - elif self.optimal_value < self.worst_feasible_value: - raise ValueError( - "The worst feasible value must be less than or equal to " - "the optimal value for maximization problems." - ) # Validate that names on optimization config are contained in names on # test function objective = self.optimization_config.objective - if isinstance(objective, MultiObjective): - objective_names = {obj.metric.name for obj in objective.objectives} - elif isinstance(objective, ScalarizedObjective): - objective_names = {metric.name for metric in objective.metrics} - else: - objective_names = {objective.metric.name} + objective_names = set(objective.metric_names) test_function_names = set(self.test_function.outcome_names) missing = objective_names - test_function_names @@ -198,7 +209,9 @@ def __post_init__(self) -> None: ) constraints = self.optimization_config.outcome_constraints - constraint_names = {c.metric.name for c in constraints} + constraint_names: set[str] = set() + for c in constraints: + constraint_names.update(c.metric_names) missing = constraint_names - test_function_names if len(missing) > 0: raise ValueError( @@ -234,9 +247,10 @@ def _get_constraints( constraint_names: Sequence[str], observe_noise_sd: bool, use_map_metric: bool = False, -) -> list[OutcomeConstraint]: +) -> tuple[list[OutcomeConstraint], list[Metric]]: """ - Create a list of ``OutcomeConstraint``s. + Create a list of ``OutcomeConstraint``s and corresponding + ``BenchmarkMetric``s. Args: constraint_names: Names of the constraints. One constraint will be @@ -246,23 +260,22 @@ def _get_constraints( where only some of the outcomes have noise levels observed. use_map_metric: Whether to use a ``BenchmarkMapMetric``. - + Returns: + A tuple of (outcome_constraints, metrics). """ metric_cls = BenchmarkMapMetric if use_map_metric else BenchmarkMetric - outcome_constraints = [ - OutcomeConstraint( - metric=metric_cls( - name=name, - lower_is_better=False, # positive slack = feasible - observe_noise_sd=observe_noise_sd, - ), - op=ComparisonOp.GEQ, - bound=0.0, - relative=False, + metrics: list[Metric] = [ + metric_cls( + name=name, + lower_is_better=False, # positive slack = feasible + observe_noise_sd=observe_noise_sd, ) for name in constraint_names ] - return outcome_constraints + outcome_constraints = [ + OutcomeConstraint(expression=f"{name} >= 0.0") for name in constraint_names + ] + return outcome_constraints, metrics def get_soo_opt_config( @@ -271,10 +284,10 @@ def get_soo_opt_config( lower_is_better: bool = True, observe_noise_sd: bool = False, use_map_metric: bool = False, -) -> OptimizationConfig: +) -> tuple[OptimizationConfig, list[Metric]]: """ Create a single-objective ``OptimizationConfig``, potentially with - constraints. + constraints, along with the corresponding ``BenchmarkMetric``s. Args: outcome_names: Names of the outcomes. If ``outcome_names`` has more than @@ -287,26 +300,29 @@ def get_soo_opt_config( observe_noise_sd: Whether the standard deviation of the observation noise is observed. Applies to all objective and constraints. use_map_metric: Whether to use a ``BenchmarkMapMetric``. + + Returns: + A tuple of (OptimizationConfig, list of BenchmarkMetrics). """ metric_cls = BenchmarkMapMetric if use_map_metric else BenchmarkMetric - objective = Objective( - metric=metric_cls( - name=outcome_names[0], - lower_is_better=lower_is_better, - observe_noise_sd=observe_noise_sd, - ), - minimize=lower_is_better, + obj_metric = metric_cls( + name=outcome_names[0], + lower_is_better=lower_is_better, + observe_noise_sd=observe_noise_sd, ) + expression = f"-{outcome_names[0]}" if lower_is_better else outcome_names[0] + objective = Objective(expression=expression) - outcome_constraints = _get_constraints( + outcome_constraints, constraint_metrics = _get_constraints( constraint_names=outcome_names[1:], observe_noise_sd=observe_noise_sd, use_map_metric=use_map_metric, ) - return OptimizationConfig( + config = OptimizationConfig( objective=objective, outcome_constraints=outcome_constraints ) + return config, [obj_metric] + constraint_metrics def get_moo_opt_config( @@ -317,9 +333,10 @@ def get_moo_opt_config( lower_is_better: bool = True, observe_noise_sd: bool = False, use_map_metric: bool = False, -) -> MultiObjectiveOptimizationConfig: +) -> tuple[MultiObjectiveOptimizationConfig, list[Metric]]: """ - Create a ``MultiObjectiveOptimizationConfig``, potentially with constraints. + Create a ``MultiObjectiveOptimizationConfig``, potentially with + constraints, along with the corresponding ``BenchmarkMetric``s. Args: outcome_names: Names of the outcomes. If ``num_constraints`` is greater @@ -338,6 +355,9 @@ def get_moo_opt_config( method, but could be. observe_noise_sd: Whether the standard deviation of the observation noise is observed. Applies to all objective and constraints. + + Returns: + A tuple of (MultiObjectiveOptimizationConfig, list of BenchmarkMetrics). """ n_objectives = len(outcome_names) - num_constraints metric_cls = BenchmarkMapMetric if use_map_metric else BenchmarkMetric @@ -349,29 +369,42 @@ def get_moo_opt_config( ) for i in range(n_objectives) ] - constraints = _get_constraints( + outcome_constraints, constraint_metrics = _get_constraints( constraint_names=outcome_names[n_objectives:], observe_noise_sd=observe_noise_sd, ) + + if n_objectives < 2: + raise ValueError( + "get_moo_opt_config requires at least 2 objectives. " + f"Got {n_objectives} objective(s) with {num_constraints} constraint(s) " + f"from {len(outcome_names)} outcome names." + ) + + # Build multi-objective expression: comma-separated, negated if lower_is_better + obj_expressions = [] + for metric in objective_metrics: + if lower_is_better: + obj_expressions.append(f"-{metric.name}") + else: + obj_expressions.append(metric.name) + objective = Objective(expression=", ".join(obj_expressions)) + + # Build objective thresholds as OutcomeConstraints + objective_thresholds = [] + for ref_p, metric in zip(ref_point, objective_metrics, strict=True): + if metric.lower_is_better: + expr = f"{metric.name} <= {ref_p}" + else: + expr = f"{metric.name} >= {ref_p}" + objective_thresholds.append(OutcomeConstraint(expression=expr)) + optimization_config = MultiObjectiveOptimizationConfig( - objective=MultiObjective( - objectives=[ - Objective(metric=metric, minimize=lower_is_better) - for metric in objective_metrics - ] - ), - objective_thresholds=[ - ObjectiveThreshold( - metric=metric, - bound=ref_p, - relative=False, - op=ComparisonOp.LEQ if metric.lower_is_better else ComparisonOp.GEQ, - ) - for ref_p, metric in zip(ref_point, objective_metrics, strict=True) - ], - outcome_constraints=constraints, + objective=objective, + objective_thresholds=objective_thresholds, + outcome_constraints=outcome_constraints, ) - return optimization_config + return optimization_config, objective_metrics + constraint_metrics def get_continuous_search_space( diff --git a/ax/benchmark/problems/surrogate/lcbench/early_stopping.py b/ax/benchmark/problems/surrogate/lcbench/early_stopping.py index 5facf2667d3..ead5ae06eab 100644 --- a/ax/benchmark/problems/surrogate/lcbench/early_stopping.py +++ b/ax/benchmark/problems/surrogate/lcbench/early_stopping.py @@ -30,7 +30,6 @@ get_lcbench_parameter_names, get_lcbench_search_space, ) -from ax.core.optimization_config import OptimizationConfig from ax.core.search_space import SearchSpace from ax.core.types import TParamValue from ax.exceptions.core import UserInputError @@ -197,7 +196,7 @@ class LearningCurveBenchmarkTestFunction(BenchmarkTestFunction): Example: test_function = LearningCurveBenchmarkTestFunction(dataset_name="vehicle") search_space = get_lcbench_search_space() - optimization_config = get_lcbench_optimization_config( + optimization_config, opt_config_metrics = get_lcbench_optimization_config( metric_name="Train/val_accuracy", observe_noise_sd=True, use_map_metric=True ) problem = BenchmarkProblem( @@ -206,6 +205,7 @@ class LearningCurveBenchmarkTestFunction(BenchmarkTestFunction): optimization_config=optimization_config, test_function=test_function, step_runtime_function=None, + opt_config_metrics=opt_config_metrics, ... ) """ @@ -292,7 +292,7 @@ def _get_lcbench_early_stopping_benchmark_problem( baseline_value = BASELINE_VALUES[dataset_name] search_space: SearchSpace = get_lcbench_search_space() - optimization_config: OptimizationConfig = get_lcbench_optimization_config( + optimization_config, opt_config_metrics = get_lcbench_optimization_config( metric_name=DEFAULT_METRIC_NAME, observe_noise_sd=False, use_map_metric=True, @@ -312,6 +312,7 @@ def _get_lcbench_early_stopping_benchmark_problem( baseline_value=baseline_value, test_function=test_function, step_runtime_function=step_runtime_function, + opt_config_metrics=opt_config_metrics, ) diff --git a/ax/benchmark/problems/surrogate/lcbench/utils.py b/ax/benchmark/problems/surrogate/lcbench/utils.py index 0cddae01975..b15da489c7a 100644 --- a/ax/benchmark/problems/surrogate/lcbench/utils.py +++ b/ax/benchmark/problems/surrogate/lcbench/utils.py @@ -6,6 +6,7 @@ # pyre-strict from ax.benchmark.benchmark_problem import get_soo_opt_config +from ax.core.metric import Metric from ax.core.optimization_config import OptimizationConfig from ax.core.parameter import ParameterType, RangeParameter from ax.core.search_space import SearchSpace @@ -116,7 +117,7 @@ def get_lcbench_optimization_config( metric_name: str = DEFAULT_METRIC_NAME, observe_noise_sd: bool = False, use_map_metric: bool = False, -) -> OptimizationConfig: +) -> tuple[OptimizationConfig, list[Metric]]: return get_soo_opt_config( outcome_names=[metric_name], lower_is_better=False, diff --git a/ax/benchmark/problems/synthetic/bandit.py b/ax/benchmark/problems/synthetic/bandit.py index 688a7dda37c..d2f7165de1d 100644 --- a/ax/benchmark/problems/synthetic/bandit.py +++ b/ax/benchmark/problems/synthetic/bandit.py @@ -47,7 +47,7 @@ def get_bandit_problem(num_choices: int = 30, num_trials: int = 3) -> BenchmarkP ) search_space = SearchSpace(parameters=[parameter]) test_function = IdentityTestFunction() - optimization_config = get_soo_opt_config( + optimization_config, opt_config_metrics = get_soo_opt_config( outcome_names=test_function.outcome_names, observe_noise_sd=True ) baselines = { @@ -74,4 +74,5 @@ def get_bandit_problem(num_choices: int = 30, num_trials: int = 3) -> BenchmarkP report_inference_value_as_trace=True, noise=GaussianNoise(noise_std=1.0), status_quo_params={"x0": num_choices // 2}, + opt_config_metrics=opt_config_metrics, ) diff --git a/ax/benchmark/problems/synthetic/discretized/mixed_integer.py b/ax/benchmark/problems/synthetic/discretized/mixed_integer.py index a6b6fe5e1d9..f2f368576d9 100644 --- a/ax/benchmark/problems/synthetic/discretized/mixed_integer.py +++ b/ax/benchmark/problems/synthetic/discretized/mixed_integer.py @@ -84,7 +84,7 @@ def _get_problem_from_common_inputs( for i in range(dim) ] ) - optimization_config = get_soo_opt_config( + optimization_config, opt_config_metrics = get_soo_opt_config( outcome_names=[metric_name], lower_is_better=lower_is_better, observe_noise_sd=observe_noise_sd, @@ -107,6 +107,7 @@ def _get_problem_from_common_inputs( num_trials=num_trials, optimal_value=optimal_value, baseline_value=baseline_value, + opt_config_metrics=opt_config_metrics, ) diff --git a/ax/benchmark/problems/synthetic/from_botorch.py b/ax/benchmark/problems/synthetic/from_botorch.py index 56aa0294db2..b41b43d9c8f 100644 --- a/ax/benchmark/problems/synthetic/from_botorch.py +++ b/ax/benchmark/problems/synthetic/from_botorch.py @@ -235,7 +235,7 @@ def create_problem_from_botorch( ) if isinstance(test_problem, MultiObjectiveTestProblem): - optimization_config = get_moo_opt_config( + optimization_config, opt_config_metrics = get_moo_opt_config( num_constraints=num_constraints, lower_is_better=lower_is_better, observe_noise_sd=observe_noise_sd, @@ -244,7 +244,7 @@ def create_problem_from_botorch( use_map_metric=use_map_metric, ) else: - optimization_config = get_soo_opt_config( + optimization_config, opt_config_metrics = get_soo_opt_config( outcome_names=test_function.outcome_names, lower_is_better=lower_is_better, observe_noise_sd=observe_noise_sd, @@ -282,6 +282,7 @@ def create_problem_from_botorch( step_runtime_function=step_runtime_function, status_quo_params=status_quo_params, auxiliary_experiments_by_purpose=auxiliary_experiments_by_purpose, + opt_config_metrics=opt_config_metrics, ) diff --git a/ax/benchmark/problems/synthetic/hss/jenatton.py b/ax/benchmark/problems/synthetic/hss/jenatton.py index ab0d7f09fc4..b20c8143c97 100644 --- a/ax/benchmark/problems/synthetic/hss/jenatton.py +++ b/ax/benchmark/problems/synthetic/hss/jenatton.py @@ -108,7 +108,7 @@ def get_jenatton_benchmark_problem( noise_std: float = 0.0, ) -> BenchmarkProblem: name = "Jenatton" + ("_observed_noise" if observe_noise_sd else "") - optimization_config = get_soo_opt_config( + optimization_config, opt_config_metrics = get_soo_opt_config( outcome_names=[name], observe_noise_sd=observe_noise_sd ) @@ -121,4 +121,5 @@ def get_jenatton_benchmark_problem( num_trials=num_trials, optimal_value=JENATTON_OPTIMAL_VALUE, baseline_value=JENATTON_BASELINE_VALUE, + opt_config_metrics=opt_config_metrics, ) diff --git a/ax/benchmark/testing/benchmark_stubs.py b/ax/benchmark/testing/benchmark_stubs.py index 39c2154c7e9..60b01e9c7dc 100644 --- a/ax/benchmark/testing/benchmark_stubs.py +++ b/ax/benchmark/testing/benchmark_stubs.py @@ -117,7 +117,7 @@ def get_soo_surrogate() -> BenchmarkProblem: experiment = get_branin_experiment(with_completed_trial=True) test_function = get_soo_surrogate_test_function() - optimization_config = get_soo_opt_config( + optimization_config, opt_config_metrics = get_soo_opt_config( outcome_names=test_function.outcome_names, observe_noise_sd=True, ) @@ -130,6 +130,7 @@ def get_soo_surrogate() -> BenchmarkProblem: optimal_value=0.0, baseline_value=3.0, test_function=test_function, + opt_config_metrics=opt_config_metrics, ) @@ -147,7 +148,7 @@ def get_moo_surrogate() -> BenchmarkProblem: test_function = SurrogateTestFunction( name="test", outcome_names=outcome_names, get_surrogate=lambda: surrogate ) - optimization_config = get_moo_opt_config( + optimization_config, opt_config_metrics = get_moo_opt_config( outcome_names=outcome_names, ref_point=[0.0, 0.0], observe_noise_sd=True, @@ -161,6 +162,7 @@ def get_moo_surrogate() -> BenchmarkProblem: optimal_value=1.0, baseline_value=0.0, test_function=test_function, + opt_config_metrics=opt_config_metrics, ) @@ -342,7 +344,7 @@ def get_async_benchmark_problem( if num_objectives == 1: # Single-objective: first outcome is objective, rest are constraints - optimization_config = get_soo_opt_config( + optimization_config, opt_config_metrics = get_soo_opt_config( outcome_names=outcome_names, lower_is_better=lower_is_better, observe_noise_sd=True, @@ -351,7 +353,7 @@ def get_async_benchmark_problem( else: # Multi-objective: pass all outcomes (objectives + constraints) # get_moo_opt_config will use the last num_constraints as constraints - optimization_config = get_moo_opt_config( + optimization_config, opt_config_metrics = get_moo_opt_config( outcome_names=outcome_names, ref_point=[1.0] * num_objectives, num_constraints=num_constraints, @@ -373,6 +375,7 @@ def get_async_benchmark_problem( else None, step_runtime_function=step_runtime_fn, report_inference_value_as_trace=report_inference_value_as_trace, + opt_config_metrics=opt_config_metrics, ) diff --git a/ax/benchmark/tests/problems/surrogate/hss/test_cifar10_surrogate.py b/ax/benchmark/tests/problems/surrogate/hss/test_cifar10_surrogate.py index 79fe78c2ecd..d08a570e124 100644 --- a/ax/benchmark/tests/problems/surrogate/hss/test_cifar10_surrogate.py +++ b/ax/benchmark/tests/problems/surrogate/hss/test_cifar10_surrogate.py @@ -18,8 +18,7 @@ def test_benchmark_creation(self) -> None: objective = benchmark.optimization_config.objective self.assertFalse(objective.minimize) - self.assertFalse(objective.metric.lower_is_better) - self.assertEqual(objective.metric.name, "CIFAR10 Test Accuracy") + self.assertEqual(objective.metric_names, ["CIFAR10 Test Accuracy"]) def test_cifar10_surrogate(self) -> None: cases = [ diff --git a/ax/benchmark/tests/problems/surrogate/hss/test_fashion_mnist_surrogate.py b/ax/benchmark/tests/problems/surrogate/hss/test_fashion_mnist_surrogate.py index 28b1fd5bf8d..2f1fb45ea18 100644 --- a/ax/benchmark/tests/problems/surrogate/hss/test_fashion_mnist_surrogate.py +++ b/ax/benchmark/tests/problems/surrogate/hss/test_fashion_mnist_surrogate.py @@ -18,8 +18,7 @@ def test_benchmark_creation(self) -> None: objective = benchmark.optimization_config.objective self.assertFalse(objective.minimize) - self.assertFalse(objective.metric.lower_is_better) - self.assertEqual(objective.metric.name, "FashionMNIST Test Accuracy") + self.assertEqual(objective.metric_names, ["FashionMNIST Test Accuracy"]) def test_fashion_mnist_surrogate(self) -> None: cases = [ diff --git a/ax/benchmark/tests/problems/surrogate/hss/test_mnist_surrogate.py b/ax/benchmark/tests/problems/surrogate/hss/test_mnist_surrogate.py index cfa00e335d2..e94c04079ec 100644 --- a/ax/benchmark/tests/problems/surrogate/hss/test_mnist_surrogate.py +++ b/ax/benchmark/tests/problems/surrogate/hss/test_mnist_surrogate.py @@ -61,8 +61,6 @@ def test_benchmark_creation(self) -> None: benchmark = get_mnist_surrogate_benchmark(num_trials=1) objective = benchmark.optimization_config.objective - metric = objective.metric - self.assertEqual(metric.name, "MNIST Test Accuracy") + self.assertEqual(objective.metric_names, ["MNIST Test Accuracy"]) self.assertFalse(objective.minimize) - self.assertFalse(metric.lower_is_better) diff --git a/ax/benchmark/tests/problems/surrogate/lcbench/test_early_stopping.py b/ax/benchmark/tests/problems/surrogate/lcbench/test_early_stopping.py index ba02cf74011..e3b38e9d865 100644 --- a/ax/benchmark/tests/problems/surrogate/lcbench/test_early_stopping.py +++ b/ax/benchmark/tests/problems/surrogate/lcbench/test_early_stopping.py @@ -63,7 +63,7 @@ def test_get_lcbench_early_stopping_problem(self) -> None: ) self.assertEqual(len(create_surrogate_regressor_call_args), 2) self.assertEqual( - problem.optimization_config.objective.metric.signature, DEFAULT_METRIC_NAME + problem.optimization_config.objective.metric_names[0], DEFAULT_METRIC_NAME ) self.assertIsNone(problem.step_runtime_function) self.assertEqual(problem.optimal_value, OPTIMAL_VALUES[dataset_name]) diff --git a/ax/benchmark/tests/problems/synthetic/hss/test_jenatton.py b/ax/benchmark/tests/problems/synthetic/hss/test_jenatton.py index bd145033419..df7b612b0ae 100644 --- a/ax/benchmark/tests/problems/synthetic/hss/test_jenatton.py +++ b/ax/benchmark/tests/problems/synthetic/hss/test_jenatton.py @@ -7,7 +7,6 @@ from random import random -from ax.benchmark.benchmark_metric import BenchmarkMetric from ax.benchmark.noise import GaussianNoise from ax.benchmark.problems.synthetic.hss.jenatton import ( get_jenatton_benchmark_problem, @@ -95,24 +94,19 @@ def test_jenatton_test_function(self) -> None: def test_create_problem(self) -> None: problem = get_jenatton_benchmark_problem() objective = problem.optimization_config.objective - metric = objective.metric - self.assertEqual(metric.name, "Jenatton") - self.assertEqual(metric.signature, "Jenatton") + self.assertEqual(objective.metric_names, ["Jenatton"]) self.assertTrue(objective.minimize) - self.assertTrue(metric.lower_is_better) self.assertEqual( assert_is_instance(problem.noise, GaussianNoise).noise_std, 0.0 ) - self.assertFalse(assert_is_instance(metric, BenchmarkMetric).observe_noise_sd) problem = get_jenatton_benchmark_problem( num_trials=10, noise_std=0.1, observe_noise_sd=True ) objective = problem.optimization_config.objective - metric = objective.metric - self.assertTrue(metric.lower_is_better) + self.assertEqual(objective.metric_names, ["Jenatton_observed_noise"]) + self.assertTrue(objective.minimize) self.assertEqual( assert_is_instance(problem.noise, GaussianNoise).noise_std, 0.1 ) - self.assertTrue(assert_is_instance(metric, BenchmarkMetric).observe_noise_sd) diff --git a/ax/benchmark/tests/problems/synthetic/test_from_botorch.py b/ax/benchmark/tests/problems/synthetic/test_from_botorch.py index 16b1f4255fd..ec3042593ad 100644 --- a/ax/benchmark/tests/problems/synthetic/test_from_botorch.py +++ b/ax/benchmark/tests/problems/synthetic/test_from_botorch.py @@ -8,7 +8,6 @@ from itertools import product import torch -from ax.benchmark.benchmark_metric import BenchmarkMapMetric, BenchmarkMetric from ax.benchmark.benchmark_problem import get_continuous_search_space from ax.benchmark.benchmark_test_functions.botorch_test import BoTorchTestFunction from ax.benchmark.noise import GaussianNoise @@ -18,7 +17,6 @@ get_augmented_branin_problem, get_augmented_branin_search_space, ) -from ax.core.objective import MultiObjective from ax.core.optimization_config import MultiObjectiveOptimizationConfig from ax.core.parameter import ChoiceParameter, RangeParameter from ax.core.types import ComparisonOp @@ -111,7 +109,7 @@ def test_single_objective_from_botorch(self) -> None: test_problem.optimal_value, botorch_test_problem._optimal_value ) # test optimization config - metric_name = test_problem.optimization_config.objective.metric.name + metric_name = test_problem.optimization_config.objective.metric_names[0] self.assertEqual(metric_name, test_problem.name) self.assertTrue(test_problem.optimization_config.objective.minimize) # test repr method @@ -123,7 +121,9 @@ def test_single_objective_from_botorch(self) -> None: outcome_constraint = ( test_problem.optimization_config.outcome_constraints[0] ) - self.assertEqual(outcome_constraint.metric.name, "constraint_slack_0") + self.assertEqual( + outcome_constraint.metric_names[0], "constraint_slack_0" + ) self.assertEqual(outcome_constraint.op, ComparisonOp.GEQ) self.assertFalse(outcome_constraint.relative) self.assertEqual(outcome_constraint.bound, 0.0) @@ -152,27 +152,20 @@ def _test_constrained_from_botorch( opt_config = ax_problem.optimization_config outcome_constraints = opt_config.outcome_constraints self.assertEqual( - [constraint.metric.name for constraint in outcome_constraints], + [constraint.metric_names[0] for constraint in outcome_constraints], [f"constraint_slack_{i}" for i in range(botorch_problem.num_constraints)], ) objective = opt_config.objective - metric = ( - objective.metrics[0] - if isinstance(objective, MultiObjective) - else objective.metric - ) - + # Verify the objective metric name matches the problem name + objective_name = objective.metric_names[0] self.assertEqual( - assert_is_instance(metric, BenchmarkMetric).observe_noise_sd, - observe_noise_sd, + objective_name, + ax_problem.name + if not isinstance(opt_config, MultiObjectiveOptimizationConfig) + else objective.metric_names[0], ) # TODO: Support observing noise variance only for some outputs - for constraint in outcome_constraints: - self.assertEqual( - assert_is_instance(constraint.metric, BenchmarkMetric).observe_noise_sd, - observe_noise_sd, - ) def test_constrained_soo_from_botorch(self) -> None: for observe_noise_sd, noise_std in product( @@ -251,12 +244,11 @@ def _test_moo_from_botorch(self, lower_is_better: bool) -> None: for t in opt_config.objective_thresholds ) ) - self.assertTrue( - all( - metric.lower_is_better is lower_is_better - for metric in opt_config.metrics.values() - ) - ) + # Verify optimization direction via objective expression + if lower_is_better: + self.assertTrue(all(w < 0 for _, w in opt_config.objective.metric_weights)) + else: + self.assertTrue(all(w > 0 for _, w in opt_config.objective.metric_weights)) def test_moo_from_botorch(self) -> None: self._test_moo_from_botorch(lower_is_better=True) @@ -361,8 +353,8 @@ def test_with_map_metric(self) -> None: num_trials=1, use_map_metric=True, ) - self.assertIsInstance( - problem.optimization_config.objective.metric, BenchmarkMapMetric + self.assertEqual( + problem.optimization_config.objective.metric_names, ["Branin"] ) self.assertEqual(problem.test_function.n_steps, 1) @@ -375,8 +367,8 @@ def test_with_map_metric(self) -> None: use_map_metric=True, n_steps=4, ) - self.assertIsInstance( - problem.optimization_config.objective.metric, BenchmarkMapMetric + self.assertEqual( + problem.optimization_config.objective.metric_names, ["Branin"] ) self.assertEqual(problem.test_function.n_steps, n_steps) @@ -387,8 +379,10 @@ def test_with_map_metric(self) -> None: num_trials=1, use_map_metric=True, ) - metric = next(iter(problem.optimization_config.metrics.values())) - self.assertIsInstance(metric, BenchmarkMapMetric) + self.assertEqual( + set(problem.optimization_config.objective.metric_names), + {"BraninCurrin_0", "BraninCurrin_1"}, + ) def test_string_test_problem_class(self) -> None: """Test that test_problem_class can be provided as a string.""" diff --git a/ax/benchmark/tests/test_benchmark.py b/ax/benchmark/tests/test_benchmark.py index 69cbea4bafe..3a97a9d569b 100644 --- a/ax/benchmark/tests/test_benchmark.py +++ b/ax/benchmark/tests/test_benchmark.py @@ -37,6 +37,7 @@ run_optimization_with_orchestrator, ) from ax.benchmark.benchmark_method import BenchmarkMethod +from ax.benchmark.benchmark_metric import BenchmarkMetric from ax.benchmark.benchmark_problem import ( BenchmarkProblem, get_continuous_search_space, @@ -218,14 +219,18 @@ def test_tracking_metrics(self) -> None: method = get_sobol_benchmark_method() problem = get_multi_objective_benchmark_problem() oc = problem.optimization_config - tracking_metric = ( - assert_is_instance(oc.objective, MultiObjective).objectives[1].metric + # Get the second objective metric name to use as a tracking metric + tracking_metric_name = oc.objective.metric_names[1] + tracking_metric = BenchmarkMetric( + name=tracking_metric_name, lower_is_better=True + ) + soo_opt_config, soo_opt_config_metrics = get_soo_opt_config( + outcome_names=[f"{problem.name}_0"], lower_is_better=True ) problem = dataclasses.replace( problem, - optimization_config=get_soo_opt_config( - outcome_names=[f"{problem.name}_0"], lower_is_better=True - ), + optimization_config=soo_opt_config, + opt_config_metrics=soo_opt_config_metrics, tracking_metrics=[tracking_metric], baseline_value=3.0, optimal_value=Branin(negate=False).optimal_value, @@ -1173,18 +1178,21 @@ def test_compute_baseline_value_from_sobol(self) -> None: test_function = IdentityTestFunction() with self.subTest("SOO, lower is better"): - opt_config = get_soo_opt_config(outcome_names=test_function.outcome_names) + opt_config, opt_config_metrics = get_soo_opt_config( + outcome_names=test_function.outcome_names + ) result = compute_baseline_value_from_sobol( optimization_config=opt_config, search_space=search_space, test_function=test_function, n_repeats=1, + opt_config_metrics=opt_config_metrics, ) self.assertEqual(result, 0) with self.subTest("SOO, Data with has_step_column=True"): map_test_function = IdentityTestFunction(n_steps=2) - map_opt_config = get_soo_opt_config( + map_opt_config, map_opt_config_metrics = get_soo_opt_config( outcome_names=test_function.outcome_names, use_map_metric=True ) result = compute_baseline_value_from_sobol( @@ -1192,11 +1200,12 @@ def test_compute_baseline_value_from_sobol(self) -> None: search_space=search_space, test_function=map_test_function, n_repeats=1, + opt_config_metrics=map_opt_config_metrics, ) self.assertEqual(result, 0) with self.subTest("SOO, higher is better"): - opt_config = get_soo_opt_config( + opt_config, opt_config_metrics = get_soo_opt_config( outcome_names=test_function.outcome_names, lower_is_better=False ) result = compute_baseline_value_from_sobol( @@ -1204,12 +1213,13 @@ def test_compute_baseline_value_from_sobol(self) -> None: search_space=search_space, test_function=test_function, n_repeats=1, + opt_config_metrics=opt_config_metrics, ) self.assertEqual(result, 4) moo_test_function = IdentityTestFunction(outcome_names=["foo", "bar"]) with self.subTest("MOO"): - moo_opt_config = get_moo_opt_config( + moo_opt_config, moo_opt_config_metrics = get_moo_opt_config( outcome_names=moo_test_function.outcome_names, ref_point=[5, 5] ) result = compute_baseline_value_from_sobol( @@ -1217,6 +1227,7 @@ def test_compute_baseline_value_from_sobol(self) -> None: search_space=search_space, test_function=moo_test_function, n_repeats=1, + opt_config_metrics=moo_opt_config_metrics, ) # (5-0) * (5-0) self.assertEqual(result, 25) @@ -1391,12 +1402,15 @@ def test_get_best_parameters(self) -> None: gs = get_sobol_generation_strategy() search_space = get_continuous_search_space(bounds=[(0, 1)]) - moo_config = get_moo_opt_config(outcome_names=["a", "b"], ref_point=[0, 0]) + moo_config, moo_metrics = get_moo_opt_config( + outcome_names=["a", "b"], ref_point=[0, 0] + ) experiment = Experiment( name="test", is_test=True, search_space=search_space, optimization_config=moo_config, + tracking_metrics=moo_metrics, ) with ( @@ -1407,7 +1421,7 @@ def test_get_best_parameters(self) -> None: ): get_best_parameters(experiment=experiment, generation_strategy=gs) - soo_config = get_soo_opt_config(outcome_names=["a"]) + soo_config, soo_metrics = get_soo_opt_config(outcome_names=["a"]) with self.subTest("Empty experiment"): result = get_best_parameters( experiment=experiment.clone_with(optimization_config=soo_config), @@ -1453,20 +1467,22 @@ def test_get_best_parameters(self) -> None: def test_worst_feasible_value_validation(self) -> None: """Test validation logic for worst_feasible_value in BenchmarkProblem.""" - search_space = get_continuous_search_space(bounds=[(0, 1), (0, 1)]) - test_function = IdentityTestFunction(outcome_names=["objective", "constraint"]) + search_space = get_continuous_search_space(bounds=[(0, 1), (0, 1), (0, 1)]) + moo_test_function = IdentityTestFunction( + outcome_names=["obj0", "obj1", "constraint"] + ) # MOO with constraints - must be 0.0 - moo_config = get_moo_opt_config( - outcome_names=["objective", "constraint"], - ref_point=[1.0], + moo_config, _moo_metrics = get_moo_opt_config( + outcome_names=["obj0", "obj1", "constraint"], + ref_point=[1.0, 1.0], num_constraints=1, ) BenchmarkProblem( name="moo_valid", optimization_config=moo_config, search_space=search_space, - test_function=test_function, + test_function=moo_test_function, num_trials=4, baseline_value=5.0, optimal_value=10.0, @@ -1477,7 +1493,7 @@ def test_worst_feasible_value_validation(self) -> None: name="moo_invalid", optimization_config=moo_config, search_space=search_space, - test_function=test_function, + test_function=moo_test_function, num_trials=4, baseline_value=5.0, optimal_value=10.0, @@ -1485,10 +1501,11 @@ def test_worst_feasible_value_validation(self) -> None: ) # SOO with constraints, `worst_feasible_value` must be provided - soo_min_config = get_soo_opt_config( + test_function = IdentityTestFunction(outcome_names=["objective", "constraint"]) + soo_min_config, _soo_min_metrics = get_soo_opt_config( outcome_names=["objective", "constraint"], lower_is_better=True ) - soo_max_config = get_soo_opt_config( + soo_max_config, _soo_max_metrics = get_soo_opt_config( outcome_names=["objective", "constraint"], lower_is_better=False ) with self.assertRaisesRegex( @@ -1551,7 +1568,9 @@ def test_worst_feasible_value_validation(self) -> None: ) # No constraints - validation skipped - no_constraint_config = get_soo_opt_config(outcome_names=["objective"]) + no_constraint_config, _no_constraint_metrics = get_soo_opt_config( + outcome_names=["objective"] + ) no_constraint_test_function = IdentityTestFunction(outcome_names=["objective"]) BenchmarkProblem( name="no_constraints", diff --git a/ax/benchmark/tests/test_benchmark_problem.py b/ax/benchmark/tests/test_benchmark_problem.py index 412f4c6c201..ee6bd72c47b 100644 --- a/ax/benchmark/tests/test_benchmark_problem.py +++ b/ax/benchmark/tests/test_benchmark_problem.py @@ -6,7 +6,6 @@ # pyre-strict -from ax.benchmark.benchmark_metric import BenchmarkMetric from ax.benchmark.benchmark_problem import ( BenchmarkProblem, get_continuous_search_space, @@ -14,18 +13,16 @@ get_soo_opt_config, ) from ax.benchmark.benchmark_test_functions.botorch_test import BoTorchTestFunction -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, OptimizationConfig, ) from ax.core.outcome_constraint import OutcomeConstraint from ax.core.search_space import SearchSpace -from ax.core.types import ComparisonOp from ax.utils.common.testutils import TestCase from botorch.test_functions.multi_objective import BraninCurrin from botorch.test_functions.synthetic import Branin -from pyre_extensions import assert_is_instance class TestBenchmarkProblem(TestCase): @@ -35,15 +32,15 @@ def setUp(self) -> None: super().setUp() def test_mismatch_of_names_on_test_function_and_opt_config_raises(self) -> None: - objectives = [ - Objective(metric=BenchmarkMetric(name, lower_is_better=True)) - for name in ["Branin", "Currin"] - ] test_function = BoTorchTestFunction( botorch_problem=Branin(), outcome_names=["Branin"] ) opt_config = MultiObjectiveOptimizationConfig( - objective=MultiObjective(objectives) + objective=Objective(expression="-Branin, -Currin"), + objective_thresholds=[ + OutcomeConstraint(expression="Branin <= 0.0"), + OutcomeConstraint(expression="Currin <= 0.0"), + ], ) with self.assertRaisesRegex( ValueError, @@ -61,14 +58,10 @@ def test_mismatch_of_names_on_test_function_and_opt_config_raises(self) -> None: baseline_value=0.0, ) - opt_config = OptimizationConfig( - objective=objectives[0], + opt_config2 = OptimizationConfig( + objective=Objective(expression="-Branin"), outcome_constraints=[ - OutcomeConstraint( - BenchmarkMetric("c", lower_is_better=False), - ComparisonOp.LEQ, - 0.0, - ) + OutcomeConstraint(expression="c <= 0.0"), ], ) with self.assertRaisesRegex( @@ -79,7 +72,7 @@ def test_mismatch_of_names_on_test_function_and_opt_config_raises(self) -> None: ): BenchmarkProblem( name="foo", - optimization_config=opt_config, + optimization_config=opt_config2, num_trials=1, optimal_value=0.0, search_space=SearchSpace(parameters=[]), @@ -89,13 +82,8 @@ def test_mismatch_of_names_on_test_function_and_opt_config_raises(self) -> None: ) def test_missing_names_on_test_function_with_scalarized_objective(self) -> None: - objective = ScalarizedObjective( - metrics=[ - BenchmarkMetric(name, lower_is_better=True) - for name in ["BraninCurrin_0", "BraninCurrin_missing"] - ], - weights=[1.0, 1.0], - minimize=True, + objective = Objective( + expression="-BraninCurrin_0 - BraninCurrin_missing", ) test_function = BoTorchTestFunction( botorch_problem=BraninCurrin(), @@ -119,57 +107,48 @@ def test_missing_names_on_test_function_with_scalarized_objective(self) -> None: ) def test_get_soo_opt_config(self) -> None: - opt_config = get_soo_opt_config( + opt_config, metrics = get_soo_opt_config( outcome_names=["foo", "bar"], lower_is_better=False, observe_noise_sd=True, ) self.assertIsInstance(opt_config.objective, Objective) self.assertEqual(len(opt_config.outcome_constraints), 1) - objective_metric = assert_is_instance( - opt_config.objective.metric, BenchmarkMetric - ) - self.assertEqual(objective_metric.name, "foo") - self.assertEqual(objective_metric.signature, "foo") - self.assertEqual(objective_metric.observe_noise_sd, True) - self.assertEqual(objective_metric.lower_is_better, False) - constraint_metric = assert_is_instance( - opt_config.outcome_constraints[0].metric, BenchmarkMetric - ) - self.assertEqual(constraint_metric.name, "bar") - self.assertEqual(constraint_metric.signature, "bar") - self.assertEqual(constraint_metric.observe_noise_sd, True) - self.assertEqual(constraint_metric.lower_is_better, False) + self.assertEqual(opt_config.objective.metric_names, ["foo"]) + self.assertFalse(opt_config.objective.minimize) + self.assertEqual(opt_config.outcome_constraints[0].metric_names, ["bar"]) + self.assertEqual(len(metrics), 2) + self.assertEqual(metrics[0].name, "foo") + self.assertEqual(metrics[1].name, "bar") def test_get_moo_opt_config(self) -> None: - opt_config = get_moo_opt_config( + opt_config, metrics = get_moo_opt_config( outcome_names=["foo", "bar", "baz", "pony"], ref_point=[3.0, 4.0], lower_is_better=False, observe_noise_sd=True, num_constraints=2, ) - self.assertEqual(len(opt_config.objective.metrics), 2) + self.assertEqual(set(opt_config.objective.metric_names), {"foo", "bar"}) self.assertEqual(len(opt_config.outcome_constraints), 2) self.assertEqual(opt_config.objective_thresholds[0].bound, 3.0) - objective_metrics = [ - assert_is_instance(metric, BenchmarkMetric) - for metric in opt_config.objective.metrics - ] - self.assertEqual(len(objective_metrics), 2) - self.assertEqual(objective_metrics[0].name, "foo") - self.assertEqual(objective_metrics[0].observe_noise_sd, True) - self.assertEqual(objective_metrics[0].lower_is_better, False) - self.assertEqual(objective_metrics[1].name, "bar") + # Check objective metric names + objective_metric_names = opt_config.objective.metric_names + self.assertEqual(len(objective_metric_names), 2) + self.assertEqual(objective_metric_names[0], "foo") + self.assertEqual(objective_metric_names[1], "bar") # Check constraints self.assertEqual(len(opt_config.outcome_constraints), 2) - constraint_metrics = [ - assert_is_instance(constraint.metric, BenchmarkMetric) - for constraint in opt_config.outcome_constraints + constraint_metric_names = [ + c.metric_names[0] for c in opt_config.outcome_constraints ] - self.assertEqual(constraint_metrics[0].name, "baz") - self.assertEqual(constraint_metrics[0].observe_noise_sd, True) - self.assertEqual(constraint_metrics[0].lower_is_better, False) + self.assertEqual(constraint_metric_names[0], "baz") + # Check metrics + self.assertEqual(len(metrics), 4) + self.assertEqual(metrics[0].name, "foo") + self.assertEqual(metrics[1].name, "bar") + self.assertEqual(metrics[2].name, "baz") + self.assertEqual(metrics[3].name, "pony") def test_get_continuous_search_space(self) -> None: bounds = [(0.0, 1.0), (2.0, 3.0)] diff --git a/ax/core/experiment.py b/ax/core/experiment.py index 9295fe4b055..4d2775f8e18 100644 --- a/ax/core/experiment.py +++ b/ax/core/experiment.py @@ -35,8 +35,11 @@ from ax.core.generator_run import GeneratorRun from ax.core.llm_provider import LLMMessage from ax.core.metric import Metric, MetricFetchE, MetricFetchResult -from ax.core.objective import MultiObjective -from ax.core.optimization_config import ObjectiveThreshold, OptimizationConfig +from ax.core.objective import Objective +from ax.core.optimization_config import ( + MultiObjectiveOptimizationConfig, + OptimizationConfig, +) from ax.core.parameter import DerivedParameter, Parameter from ax.core.runner import Runner from ax.core.search_space import SearchSpace @@ -93,7 +96,7 @@ def __init__( search_space: SearchSpace, name: str | None = None, optimization_config: OptimizationConfig | None = None, - tracking_metrics: list[Metric] | None = None, + metrics: Sequence[Metric] | None = None, runner: Runner | None = None, status_quo: Arm | None = None, description: str | None = None, @@ -104,6 +107,7 @@ def __init__( auxiliary_experiments_by_purpose: None | (dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]]) = None, default_trial_type: str | None = None, + tracking_metrics: list[Metric] | None = None, ) -> None: """Inits Experiment. @@ -111,7 +115,6 @@ def __init__( search_space: Search space of the experiment. name: Name of the experiment. optimization_config: Optimization config of the experiment. - tracking_metrics: Additional tracking metrics not used for optimization. runner: Default runner used for trials on this experiment. status_quo: Arm representing existing "control" arm. description: Description of the experiment. @@ -126,6 +129,7 @@ def __init__( default_data_type: Deprecated and ignored. auxiliary_experiments_by_purpose: Dictionary of auxiliary experiments for different purposes (e.g., transfer learning). + tracking_metrics: Deprecated. Use ``metrics`` directly instead. """ if default_data_type is not None: warnings.warn( @@ -145,7 +149,7 @@ def __init__( self.data: Data = Data() self._experiment_type: str | None = experiment_type self._optimization_config: OptimizationConfig | None = None - self._tracking_metrics: dict[str, Metric] = {} + self._metrics: dict[str, Metric] = {} self._time_created: datetime = datetime.now() self._status: ExperimentStatus | None = None self._trials: dict[int, BaseTrial] = {} @@ -188,11 +192,28 @@ def __init__( ) in self._initial_auxiliary_experiments_by_purpose.items() } - self.add_tracking_metrics(tracking_metrics or []) + # Attach metrics from both tracking_metrics and metrics, preferring metrics if + # a naming collision occurs. + for m in [*(tracking_metrics or []), *(metrics or [])]: + self._metrics[m.name] = m # call setters defined below self.status_quo = status_quo if optimization_config is not None: + # Auto-register placeholder Metric objects for any optimization + # config metric names not already on the experiment. The + # ``optimization_config`` setter (below) validates that every + # referenced metric name exists in ``self._metrics``, so + # placeholders are needed here to satisfy that check. These + # placeholders are later replaced in two ways: + # 1. During SQA decoding, the decoder overwrites them with + # properly typed metrics from the database + # (see ``sqa_store/decoder.py``). + # 2. Any call to ``add_metric`` or ``update_metric`` with a + # real Metric of the same name will overwrite a placeholder. + for name in optimization_config.metric_names: + if name not in self._metrics: + self._metrics[name] = Metric(name=name) self.optimization_config = optimization_config # Keyed on tuple[trial_index, metric_name]. @@ -549,21 +570,13 @@ def optimization_config(self, optimization_config: OptimizationConfig) -> None: f"`{Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF.value}` " "property that is set to `True` on this experiment." ) - for metric_name in optimization_config.metrics.keys(): - if metric_name in self._tracking_metrics: - self.remove_tracking_metric(metric_name) - # add metrics from the previous optimization config that are not in the new - # optimization config as tracking metrics - prev_optimization_config = self._optimization_config + for metric_name in optimization_config.metric_names: + if metric_name not in self._metrics: + raise ValueError( + f"Metric '{metric_name}' referenced in optimization config " + "but not found on experiment. Add it first with add_metric()." + ) self._optimization_config = optimization_config - if prev_optimization_config is not None: - metrics_to_track = ( - set(prev_optimization_config.metrics.keys()) - - set(optimization_config.metrics.keys()) - - {Keys.DEFAULT_OBJECTIVE_NAME.value} # remove default objective - ) - for metric_name in metrics_to_track: - self.add_tracking_metric(prev_optimization_config.metrics[metric_name]) @property def is_moo_problem(self) -> bool: @@ -702,20 +715,14 @@ def _format_experiment_summary_for_llm(self) -> str: opt_config = self.optimization_config if opt_config is not None: opt_lines: list[str] = ["### Optimization Config"] + objective = opt_config.objective - if isinstance(objective, MultiObjective): - for obj in objective.objectives: - direction = "minimize" if obj.minimize else "maximize" - opt_lines.append(f"- Objective: {direction} `{obj.metric.name}`") - else: - direction = "minimize" if objective.minimize else "maximize" - opt_lines.append(f"- Objective: {direction} `{objective.metric.name}`") + for metric_name, weight in objective.metric_weights: + direction = "minimize" if weight < 0 else "maximize" + opt_lines.append(f"- Objective: {direction} `{metric_name}`") + for constraint in opt_config.outcome_constraints: - op_str = "<=" if constraint.op == ComparisonOp.LEQ else ">=" - opt_lines.append( - f"- Constraint: `{constraint.metric.name}` " - f"{op_str} {constraint.bound}" - ) + opt_lines.append(f"- Constraint: `{constraint.expression}` ") sections.append("\n".join(opt_lines)) # Trial data @@ -790,77 +797,131 @@ def _format_experiment_summary_for_llm(self) -> str: @property def tracking_metrics(self) -> list[Metric]: - return list(self._tracking_metrics.values()) + """Metrics that are tracked but not part of the optimization config.""" + if self._optimization_config is None: + return list(self._metrics.values()) + opt_metric_names = self._optimization_config.metric_names + return [m for name, m in self._metrics.items() if name not in opt_metric_names] - def add_tracking_metric(self, metric: Metric) -> Self: + @property + def optimization_config_metrics(self) -> list[Metric]: + """Metrics that are part of the optimization config.""" + if self._optimization_config is None: + return [] + opt_metric_names = self._optimization_config.metric_names + return self.get_metrics(metric_names=list(opt_metric_names)) + + def get_metric(self, name: str) -> Metric: + """Get a single Metric by name. + + Args: + name: The name of the metric to retrieve. + + Returns: + The Metric object with the given name. + + Raises: + KeyError: If no metric with the given name exists. + """ + if name not in self._metrics: + raise KeyError( + f"Metric '{name}' not found. Available: {list(self._metrics.keys())}" + ) + return self._metrics[name] + + def add_metric(self, metric: Metric) -> Self: """Add a new metric to the experiment. + Metrics that are not referenced by the experiment's optimization config + (i.e. not part of an objective or outcome constraint) are automatically + treated as tracking metrics. + Args: metric: Metric to be added. """ - if metric.name in self._tracking_metrics: + if metric.name in self._metrics: raise ValueError( f"Metric `{metric.name}` already defined on experiment. " - "Use `update_tracking_metric` to update an existing metric definition." + "Use `update_metric` to update an existing metric definition." ) - - optimization_config = self.optimization_config - if optimization_config and metric.name in optimization_config.metrics: - raise ValueError( - f"Metric `{metric.name}` already present in experiment's " - "OptimizationConfig. Set a new OptimizationConfig without this metric " - "before adding it to tracking metrics." - ) - - self._tracking_metrics[metric.name] = metric + self._metrics[metric.name] = metric return self - def add_tracking_metrics(self, metrics: list[Metric]) -> Self: - """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 - - Args: - metrics: Metrics to be added. - """ - # Before setting any metrics, we validate none are already on - # the experiment + def add_tracking_metric(self, metric: Metric) -> Self: + """*Deprecated.* Use ``add_metric`` instead.""" + warnings.warn( + "add_tracking_metric is deprecated. Use add_metric instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.add_metric(metric) + + def add_tracking_metrics(self, metrics: list[Metric]) -> Experiment: + """*Deprecated.* Use ``add_metric`` instead.""" + warnings.warn( + "add_tracking_metrics is deprecated. Use add_metric instead.", + DeprecationWarning, + stacklevel=2, + ) for metric in metrics: - self.add_tracking_metric(metric) + self.add_metric(metric) return self - def update_tracking_metric(self, metric: Metric) -> Self: + def update_metric(self, metric: Metric) -> Self: """Redefine a metric that already exists on the experiment. Args: metric: New metric definition. """ - if metric.name not in self._tracking_metrics: + if metric.name not in self._metrics: raise ValueError(f"Metric `{metric.name}` doesn't exist on experiment.") - - self._tracking_metrics[metric.name] = metric + self._metrics[metric.name] = metric return self - def remove_tracking_metric(self, metric_name: str) -> Self: - """Remove a metric that already exists on the experiment. + def update_tracking_metric(self, metric: Metric) -> Experiment: + """*Deprecated.* Use ``update_metric`` instead.""" + warnings.warn( + "update_tracking_metric is deprecated. Use update_metric instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.update_metric(metric) + + def remove_metric(self, metric_name: str) -> Self: + """Remove a metric from the experiment. Args: metric_name: Unique name of metric to remove. + + Raises: + ValueError: If the metric is referenced by the optimization config. """ - if metric_name not in self._tracking_metrics: + if metric_name not in self._metrics: raise ValueError(f"Metric `{metric_name}` doesn't exist on experiment.") - - del self._tracking_metrics[metric_name] + if ( + self._optimization_config is not None + and metric_name in self._optimization_config.metric_names + ): + raise ValueError( + f"Metric `{metric_name}` is referenced by the optimization config " + "and cannot be removed. Update the optimization config first." + ) + del self._metrics[metric_name] return self + def remove_tracking_metric(self, metric_name: str) -> Experiment: + """*Deprecated.* Use ``remove_metric`` instead.""" + warnings.warn( + "remove_tracking_metric is deprecated. Use remove_metric instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.remove_metric(metric_name) + @property def metrics(self) -> dict[str, Metric]: """The metrics attached to the experiment.""" - optimization_config_metrics: dict[str, Metric] = {} - if self.optimization_config is not None: - optimization_config_metrics = self.optimization_config.metrics - return {**self._tracking_metrics, **optimization_config_metrics} + return dict(self._metrics) @property def signature_to_metric(self) -> dict[str, Metric]: @@ -953,7 +1014,7 @@ def bulk_configure_metrics_of_class( new_metric = metric for param_name, param_value in attributes_to_update.items(): setattr(new_metric, param_name, param_value) - self.metrics[metric.name] = new_metric + self._metrics[metric.name] = new_metric return def fetch_trials_data_results( @@ -1376,7 +1437,7 @@ def trial_indices_with_data( "Cannot find trials with data for optimization config metrics " "because no optimization config has been defined." ) - metric_names = set(self.optimization_config.metrics.keys()) + metric_names = self.optimization_config.metric_names else: metric_names = set(self.metrics.keys()) if len(metric_names) == 0: @@ -2053,8 +2114,8 @@ def clone_with( else optimization_config ) tracking_metrics = ( - [m.clone() for m in self.tracking_metrics] - if (tracking_metrics is None and self.tracking_metrics is not None) + [m.clone() for m in self.metrics.values()] + if tracking_metrics is None else tracking_metrics ) @@ -2150,28 +2211,36 @@ def metric_config_summary_df(self) -> pd.DataFrame: records[m.name] = m.summary_dict if self.optimization_config is not None: opt_config = self.optimization_config - if self.is_moo_problem: - multi_objective = assert_is_instance( - opt_config.objective, MultiObjective - ) - objectives = multi_objective.objectives + objective = opt_config.objective + if objective.is_multi_objective: + parts = [p.strip() for p in objective.expression.split(",")] + sub_objectives = [Objective(expression=part) for part in parts] else: - objectives = [opt_config.objective] - for objective in objectives: - records[objective.metric.name][METRIC_DF_COLNAMES["goal"]] = ( - "minimize" if objective.minimize else "maximize" - ) + sub_objectives = [objective] + for sub_obj in sub_objectives: + obj_name = sub_obj.metric_names[0] + if obj_name in records: + records[obj_name][METRIC_DF_COLNAMES["goal"]] = ( + "minimize" if sub_obj.minimize else "maximize" + ) + + objective_threshold_names: set[str] = set() + if isinstance(opt_config, MultiObjectiveOptimizationConfig): + objective_threshold_names = { + ot.metric_names[0] for ot in opt_config.objective_thresholds + } for constraint in opt_config.all_constraints: - if not isinstance(constraint, ObjectiveThreshold): - records[constraint.metric.name][METRIC_DF_COLNAMES["goal"]] = ( - "constrain" - ) + c_name = constraint.metric_names[0] + if c_name not in objective_threshold_names: + if c_name in records: + records[c_name][METRIC_DF_COLNAMES["goal"]] = "constrain" op = ">= " if constraint.op == ComparisonOp.GEQ else "<= " relative = "%" if constraint.relative else "" - records[constraint.metric.name][METRIC_DF_COLNAMES["bound"]] = ( - f"{op}{constraint.bound}{relative}" - ) + if c_name in records: + records[c_name][METRIC_DF_COLNAMES["bound"]] = ( + f"{op}{constraint.bound}{relative}" + ) for metric in self.tracking_metrics or []: records[metric.name][METRIC_DF_COLNAMES["goal"]] = "track" diff --git a/ax/core/multi_type_experiment.py b/ax/core/multi_type_experiment.py index df4ff76525a..f6a1ce91add 100644 --- a/ax/core/multi_type_experiment.py +++ b/ax/core/multi_type_experiment.py @@ -96,6 +96,16 @@ def __init__( default_data_type=default_data_type, ) + # Ensure tracking metrics are registered in _metric_to_trial_type. + # super().__init__ sets self._metrics directly, bypassing + # add_tracking_metric, so tracking metrics won't be in + # _metric_to_trial_type yet. + for m in tracking_metrics or []: + if m.name not in self._metric_to_trial_type: + self._metric_to_trial_type[m.name] = none_throws( + self._default_trial_type + ) + def add_trial_type(self, trial_type: str, runner: Runner) -> Self: """Add a new trial_type to be supported by this experiment. @@ -116,7 +126,7 @@ def add_trial_type(self, trial_type: str, runner: Runner) -> Self: def optimization_config(self, optimization_config: OptimizationConfig) -> None: # pyre-fixme[16]: `Optional` has no attribute `fset`. Experiment.optimization_config.fset(self, optimization_config) - for metric_name in optimization_config.metrics.keys(): + for metric_name in optimization_config.metric_names: # Optimization config metrics are required to be the default trial type # currently. TODO: remove that restriction (T202797235) self._metric_to_trial_type[metric_name] = none_throws( @@ -221,8 +231,8 @@ def update_tracking_metric( metric.name, self._default_trial_type ) oc = self.optimization_config - oc_metrics = oc.metrics if oc else [] - if metric.name in oc_metrics and trial_type != self._default_trial_type: + oc_metric_names = oc.metric_names if oc else set() + if metric.name in oc_metric_names and trial_type != self._default_trial_type: raise ValueError( f"Metric `{metric.name}` must remain a " f"`{self._default_trial_type}` metric because it is part of the " @@ -239,11 +249,11 @@ def update_tracking_metric( @copy_doc(Experiment.remove_tracking_metric) def remove_tracking_metric(self, metric_name: str) -> Self: - if metric_name not in self._tracking_metrics: + if metric_name not in self._metrics: raise ValueError(f"Metric `{metric_name}` doesn't exist on experiment.") # Required fields - del self._tracking_metrics[metric_name] + del self._metrics[metric_name] del self._metric_to_trial_type[metric_name] # Optional @@ -302,7 +312,7 @@ def metric_to_trial_type(self) -> dict[str, str]: """ opt_config_types = { metric_name: self.default_trial_type - for metric_name in self.optimization_config.metrics.keys() + for metric_name in self.optimization_config.metric_names } return {**opt_config_types, **self._metric_to_trial_type} diff --git a/ax/core/objective.py b/ax/core/objective.py index 8036ff57ef8..1ae59d8b611 100644 --- a/ax/core/objective.py +++ b/ax/core/objective.py @@ -8,83 +8,189 @@ from __future__ import annotations -from collections.abc import Iterable +import warnings from typing import Self from ax.core.metric import Metric from ax.exceptions.core import UserInputError from ax.utils.common.base import SortableBase -from pyre_extensions import none_throws +from ax.utils.common.sympy import ( + extract_metric_names_from_objective_expr, + extract_metric_weights_from_objective_expr, + parse_objective_expression, +) class Objective(SortableBase): - """Base class for representing an objective. + """Class for representing an optimization objective via an expression string. + + An objective is defined by an expression string that describes the quantity + to be **maximized**. To minimize a metric, negate it in the expression + (e.g. ``"-loss"``). + + Examples:: + + Objective(expression="accuracy") # maximize accuracy + Objective(expression="-loss") # minimize loss + Objective(expression="2*acc + recall") # scalarized objective + Objective(expression="acc, -loss") # multi-objective Attributes: - minimize: If True, minimize metric. + expression: The expression string defining this objective. """ - def __init__(self, metric: Metric, minimize: bool | None = None) -> None: + def __init__( + self, + expression: str | None = None, + *, + # Deprecated backward-compat kwargs + metric: Metric | None = None, + minimize: bool | None = None, + ) -> None: """Create a new objective. Args: - metric: The metric to be optimized. - minimize: If True, minimize metric. If None, will be set based on the - `lower_is_better` property of the metric (if that is not specified, - will raise a `UserInputError`). - + expression: A string expression defining the objective. Metrics are + referenced by name. The expression is always interpreted as + something to be **maximized**. Use negation to minimize (e.g. + ``"-loss"``). Comma-separated expressions define a + multi-objective (e.g. ``"accuracy, -loss"``). + metric: *Deprecated.* A single metric to optimize. Provide + ``expression`` instead. + minimize: *Deprecated.* Only used together with ``metric``. If + ``True`` the expression is negated so the metric is minimized. """ - lower_is_better = metric.lower_is_better - if minimize is None: - if lower_is_better is None: + if metric is not None: + # Backward-compat path + warnings.warn( + "Passing `metric` to Objective is deprecated. " + "Use `expression` instead, e.g. " + f'Objective(expression="{metric.name}") or ' + f'Objective(expression="-{metric.name}").', + DeprecationWarning, + stacklevel=2, + ) + if expression: + raise UserInputError("Cannot specify both `expression` and `metric`.") + # Resolve minimize + lower_is_better = metric.lower_is_better + if minimize is None: + if lower_is_better is None: + raise UserInputError( + f"Metric {metric.name} does not specify `lower_is_better` " + "and `minimize` is not specified. At least one of these " + "must be specified." + ) + else: + minimize = lower_is_better + elif lower_is_better is not None and lower_is_better != minimize: raise UserInputError( - f"Metric {metric.name} does not specify `lower_is_better` " - "and `minimize` is not specified. At least one of these " - "must be specified." + f"Metric {metric.name} specifies {lower_is_better=}, " + "which doesn't match the specified optimization direction " + f"{minimize=}." ) - else: - minimize = lower_is_better - elif lower_is_better is not None and lower_is_better != minimize: + expression = f"-{metric.name}" if minimize else metric.name + + if expression is None: raise UserInputError( - f"Metric {metric.name} specifies {lower_is_better=}, " - "which doesn't match the specified optimization direction " - f"{minimize=}." + "An expression string is required. " + "Example: Objective(expression='accuracy') or " + "Objective(expression='-loss')." ) - self._metric: Metric = metric - self.minimize: bool = none_throws(minimize) - @property - def metric(self) -> Metric: - """Get the objective metric.""" - return self._metric + self._expression_str: str = expression + + # Eagerly validate: error on duplicate metric names + parsed = parse_objective_expression(expression) + sub_exprs = parsed if isinstance(parsed, tuple) else (parsed,) + seen: list[str] = [] + for sub_expr in sub_exprs: + for name in extract_metric_names_from_objective_expr(sub_expr): + if name in seen: + raise UserInputError( + f"Metric '{name}' appears more than once in the objective " + f"expression '{expression}'." + ) + seen.append(name) @property - def metrics(self) -> list[Metric]: - """Get a list of objective metrics.""" - return [self._metric] + def expression(self) -> str: + """Get the expression string defining this objective.""" + return self._expression_str @property def metric_names(self) -> list[str]: - """Get a list of objective metric names.""" - return [m.name for m in self.metrics] + """Get a list of all metric names referenced in the expression.""" + parsed = parse_objective_expression(self._expression_str) + sub_exprs = parsed if isinstance(parsed, tuple) else (parsed,) + names: list[str] = [] + for sub_expr in sub_exprs: + for name in extract_metric_names_from_objective_expr(sub_expr): + if name not in names: + names.append(name) + return names + + @property + def metric_weights(self) -> list[tuple[str, float]]: + """Get a list of (metric_name, weight) tuples in the expression.""" + parsed = parse_objective_expression(self._expression_str) + sub_exprs = parsed if isinstance(parsed, tuple) else (parsed,) + + result: list[tuple[str, float]] = [] + for sub_expr in sub_exprs: + result.extend(extract_metric_weights_from_objective_expr(sub_expr)) + + return result + + @property + def is_multi_objective(self) -> bool: + """True if the objective has multiple comma-separated + sub-objectives.""" + return isinstance(parse_objective_expression(self._expression_str), tuple) + + @property + def is_scalarized_objective(self) -> bool: + """True if the objective is a single linear combination of + multiple metrics.""" + return not self.is_multi_objective and len(self.metric_names) > 1 @property - def metric_signatures(self) -> list[str]: - """Get a list of objective metric signatures.""" - return [m.signature for m in self.metrics] + def is_single_objective(self) -> bool: + """True if the objective has exactly one metric and is not multi.""" + return not self.is_multi_objective and not self.is_scalarized_objective + + @property + def minimize(self) -> bool: + """Whether this objective is minimizing. + + Only meaningful for single-metric objectives. A negative coefficient + in the expression indicates minimization. + + Raises: + UserInputError: If the objective is multi or scalarized. + """ + if self.is_multi_objective or self.is_scalarized_objective: + raise UserInputError( + "`minimize` is only defined for single-metric objectives. " + "For scalarized or multi-objectives, inspect the expression " + "or metric_weights directly." + ) + weights = self.metric_weights + if len(weights) == 1: + return weights[0][1] < 0 + return False + + def get_unconstrainable_metric_names(self) -> list[str]: + """Return metric names that are incompatible with + OutcomeConstraints.""" + return self.metric_names def clone(self) -> Self: """Create a copy of the objective.""" - return self.__class__(self.metric.clone(), self.minimize) + return self.__class__(expression=self._expression_str) def __repr__(self) -> str: - return 'Objective(metric_name="{}", minimize={})'.format( - self.metric.name, self.minimize - ) - - def get_unconstrainable_metrics(self) -> list[Metric]: - """Return a list of metrics that are incompatible with OutcomeConstraints.""" - return self.metrics + return f'Objective(expression="{self._expression_str}")' @property def _unique_id(self) -> str: @@ -92,61 +198,52 @@ def _unique_id(self) -> str: class MultiObjective(Objective): - """Class for an objective composed of a multiple component objectives. - - The Acquisition function determines how the objectives are weighted. - - Attributes: - objectives: List of objectives. + """*Deprecated.* Use ``Objective`` with a comma-separated expression + instead. + + Example:: + + # Old + MultiObjective(objectives=[ + Objective(expression="acc"), + Objective(expression="-loss"), + ]) + # New + Objective(expression="acc, -loss") """ def __init__(self, objectives: list[Objective]) -> None: - """Create a new objective. - - Args: - objectives: The list of objectives to be jointly optimized. - - """ - self._objectives: list[Objective] = objectives + warnings.warn( + "MultiObjective is deprecated. Use Objective with a " + "comma-separated expression instead, e.g. " + 'Objective(expression="acc, -loss").', + DeprecationWarning, + stacklevel=2, + ) if any(isinstance(o, ScalarizedObjective) for o in objectives): raise NotImplementedError( "Scalarized objectives are not supported for a `MultiObjective`." ) + expression = ", ".join(obj.expression for obj in objectives) + super().__init__(expression=expression) - @property - def metric(self) -> Metric: - """Override base method to error.""" - raise NotImplementedError( - f"{type(self).__name__} is composed of multiple metrics" - ) - - @property - def metrics(self) -> list[Metric]: - """Get the objective metrics.""" - return [o.metric for o in self._objectives] - - @property - def objectives(self) -> list[Objective]: - """Get the objectives.""" - return self._objectives - - def clone(self) -> Self: - """Create a copy of the objective.""" - return self.__class__(objectives=[o.clone() for o in self.objectives]) - - def __repr__(self) -> str: - return f"MultiObjective(objectives={self.objectives})" + def clone(self) -> Objective: # pyre-ignore[15]: Inconsistent override + """Clone as a base Objective (MultiObjective is deprecated).""" + return Objective(expression=self._expression_str) class ScalarizedObjective(Objective): - """Class for an objective composed of a linear scalarization of metrics. + """*Deprecated.* Use ``Objective`` with a weighted expression instead. - Attributes: - metrics: List of metrics. - weights: Weights for scalarization; default to 1. - """ + Example:: - weights: list[float] + # Old + ScalarizedObjective( + metrics=[m1, m2], weights=[2.0, 1.0], minimize=True + ) + # New + Objective(expression="-2*m1 - m2") + """ def __init__( self, @@ -154,81 +251,46 @@ def __init__( weights: list[float] | None = None, minimize: bool = False, ) -> None: - """Create a new objective. - - Args: - metric: The metric to be optimized. - weights: The weights for the linear combination of metrics. - minimize: If true, minimize the linear combination. - - """ + warnings.warn( + "ScalarizedObjective is deprecated. Use Objective with a " + "weighted expression instead, e.g. " + 'Objective(expression="2*m1 + m2").', + DeprecationWarning, + stacklevel=2, + ) if weights is None: weights = [1.0 for _ in metrics] else: if len(weights) != len(metrics): raise ValueError("Length of weights must equal length of metrics") - # Check if the optimization direction is consistent with - # `lower_is_better` (if specified). + # Validate lower_is_better consistency for m, w in zip(metrics, weights): is_minimized = minimize if w > 0 else not minimize if m.lower_is_better is not None and is_minimized != m.lower_is_better: raise ValueError( f"Metric with name {m.name} specifies `lower_is_better` = " f"{m.lower_is_better}, which doesn't match the specified " - "optimization direction. You most likely want to flip the sign of " - "the corresponding metric weight." + "optimization direction. You most likely want to flip the " + "sign of the corresponding metric weight." ) - self._metrics = metrics - self.weights = weights - self.minimize = minimize - - @property - def metric(self) -> Metric: - """Override base method to error.""" - raise NotImplementedError( - f"{type(self).__name__} is composed of multiple metrics" - ) - - @property - def metrics(self) -> list[Metric]: - """Get the metrics.""" - return self._metrics - - @property - def metric_weights(self) -> Iterable[tuple[Metric, float]]: - """Get the metrics and weights.""" - return zip(self.metrics, self.weights) - - @property - def expression(self) -> str: - """ - Get a human-readable mathematical expression of the scalarized objective. - - Returns a string representation showing the weighted combination of metrics, - e.g., "m1 + 2.0*m3" or "accuracy - 0.5*latency". - """ - parts = [] - for metric, weight in self.metric_weights: - if weight == 1.0: - parts.append(metric.name) - elif weight == -1.0: - parts.append(f"-{metric.name}") + # Build expression string + # When minimize=True, flip sign so expression represents maximization + sign = -1.0 if minimize else 1.0 + parts: list[str] = [] + for m, w in zip(metrics, weights): + effective_w = sign * w + if effective_w == 1.0: + parts.append(m.name) + elif effective_w == -1.0: + parts.append(f"-{m.name}") else: - parts.append(f"{weight}*{metric.name}") - - return " + ".join(parts).replace(" + -", " - ") + parts.append(f"{effective_w}*{m.name}") - def clone(self) -> Self: - """Create a copy of the objective.""" - return self.__class__( - metrics=[m.clone() for m in self.metrics], - weights=self.weights.copy(), - minimize=self.minimize, - ) + expression_str = " + ".join(parts).replace(" + -", " - ") + super().__init__(expression=expression_str) - def __repr__(self) -> str: - return "ScalarizedObjective(metric_names={}, weights={}, minimize={})".format( - [metric.name for metric in self.metrics], self.weights, self.minimize - ) + def clone(self) -> Objective: # pyre-ignore[15]: Inconsistent override + """Clone as a base Objective (ScalarizedObjective is deprecated).""" + return Objective(expression=self._expression_str) diff --git a/ax/core/optimization_config.py b/ax/core/optimization_config.py index ba5d1138f95..42cfdaf0574 100644 --- a/ax/core/optimization_config.py +++ b/ax/core/optimization_config.py @@ -12,27 +12,24 @@ from typing import Self from ax.core.arm import Arm -from ax.core.metric import Metric -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.outcome_constraint import ( ComparisonOp, ObjectiveThreshold, OutcomeConstraint, - ScalarizedOutcomeConstraint, ) from ax.exceptions.core import UserInputError from ax.utils.common.base import Base -from pyre_extensions import assert_is_instance -TRefPoint = list[ObjectiveThreshold] +TRefPoint = list[OutcomeConstraint] # Sentinels for default arguments when None is a valid input -_NO_OUTCOME_CONSTRAINTS = [ - OutcomeConstraint(Metric("placeholder", lower_is_better=True), ComparisonOp.GEQ, 0) +_NO_OUTCOME_CONSTRAINTS: list[OutcomeConstraint] = [ + OutcomeConstraint(expression="__placeholder__ >= 0") ] -_NO_OBJECTIVE_THRESHOLDS = [ - ObjectiveThreshold(Metric("placeholder", lower_is_better=True), 0) +_NO_OBJECTIVE_THRESHOLDS: list[OutcomeConstraint] = [ + OutcomeConstraint(expression="__placeholder__ >= 0") ] _NO_PRUNING_TARGET_PARAMETERIZATION = Arm(parameters={}) @@ -135,29 +132,25 @@ def outcome_constraints(self) -> list[OutcomeConstraint]: return self._outcome_constraints @property - def metrics(self) -> dict[str, Metric]: - """Returns mapping of name to metric.""" - constraint_metrics = { - oc.metric.name: oc.metric - for oc in self.all_constraints - if not isinstance(oc, ScalarizedOutcomeConstraint) - } - scalarized_constraint_metrics = { - metric.name: metric - for oc in self.all_constraints - if isinstance(oc, ScalarizedOutcomeConstraint) - for metric in oc.metrics - } - objective_metrics = {metric.name: metric for metric in self.objective.metrics} - return { - **constraint_metrics, - **scalarized_constraint_metrics, - **objective_metrics, - } + def objective_thresholds(self) -> list[OutcomeConstraint]: + """Get objective thresholds.""" + return [ + threshold + for threshold in self.outcome_constraints + if threshold.metric_names[0] in self.objective.metric_names + ] + + @property + def metric_names(self) -> set[str]: + """All metric names referenced by the objective and constraints.""" + names: set[str] = set(self.objective.metric_names) + for oc in self.all_constraints: + names.update(oc.metric_names) + return names @property def is_moo_problem(self) -> bool: - return self.objective is not None and isinstance(self.objective, MultiObjective) + return self.objective is not None and self.objective.is_multi_objective @property def is_bope_problem(self) -> bool: @@ -195,55 +188,47 @@ def _validate_transformed_optimization_config( objective: Metric+direction to use for the optimization. outcome_constraints: Constraints to validate. """ - if type(objective) is MultiObjective: - # Raise error on exact equality; `ScalarizedObjective` is OK + if objective.is_multi_objective: + # Raise error on multi-objective; `ScalarizedObjective` is OK raise ValueError( "OptimizationConfig does not support MultiObjective. " "Use MultiObjectiveOptimizationConfig instead." ) outcome_constraints = outcome_constraints or [] - unconstrainable_metrics = objective.get_unconstrainable_metrics() + unconstrainable_metric_names = objective.get_unconstrainable_metric_names() OptimizationConfig._validate_outcome_constraints( - unconstrainable_metrics=unconstrainable_metrics, + unconstrainable_metric_names=unconstrainable_metric_names, outcome_constraints=outcome_constraints, ) @staticmethod def _validate_outcome_constraints( - unconstrainable_metrics: list[Metric], + unconstrainable_metric_names: list[str], outcome_constraints: list[OutcomeConstraint], ) -> None: - constraint_metric_map = {} + # Build a set of all metric names referenced in constraints + constraint_metric_names: set[str] = set() for oc in outcome_constraints: - if isinstance(oc, ScalarizedOutcomeConstraint): - for m in oc.metrics: - constraint_metric_map[m.signature] = m.name - else: - constraint_metric_map[oc.metric.signature] = oc.metric.name - - for metric in unconstrainable_metrics: - if metric.signature in constraint_metric_map: + constraint_metric_names.update(oc.metric_names) + + for name in unconstrainable_metric_names: + if name in constraint_metric_names: raise ValueError("Cannot constrain on objective metric.") def constraint_key(oc: OutcomeConstraint) -> str: - return ( - str(oc) - if isinstance(oc, ScalarizedOutcomeConstraint) - else oc.metric.signature - ) + # For multi-metric constraints, use the full string representation + # For single-metric, use the metric name + if len(oc.metric_names) > 1: + return str(oc) + return oc.metric_names[0] sorted_constraints = sorted(outcome_constraints, key=constraint_key) - for metric_signature, constraints_itr in groupby( - sorted_constraints, constraint_key - ): + for key, constraints_itr in groupby(sorted_constraints, constraint_key): constraints: list[OutcomeConstraint] = list(constraints_itr) constraints_len = len(constraints) if constraints_len == 2: if constraints[0].op == constraints[1].op: - raise ValueError( - f"Duplicate outcome constraints " - f"{constraint_metric_map[metric_signature]}" - ) + raise ValueError(f"Duplicate outcome constraints {key}") lower_bound_idx = 0 if constraints[0].op == ComparisonOp.GEQ else 1 upper_bound_idx = 1 - lower_bound_idx lower_bound = constraints[lower_bound_idx].bound @@ -251,14 +236,10 @@ def constraint_key(oc: OutcomeConstraint) -> str: if lower_bound >= upper_bound: raise ValueError( f"Lower bound {lower_bound} is >= upper bound " - f"{upper_bound} for " - f"{constraint_metric_map[metric_signature]}" + f"{upper_bound} for {key}" ) elif constraints_len > 2: - raise ValueError( - "Duplicate outcome constraints " - f"{constraint_metric_map[metric_signature]}" - ) + raise ValueError(f"Duplicate outcome constraints {key}") def __repr__(self) -> str: return ( @@ -288,9 +269,9 @@ class MultiObjectiveOptimizationConfig(OptimizationConfig): def __init__( self, - objective: MultiObjective | ScalarizedObjective, + objective: Objective, outcome_constraints: list[OutcomeConstraint] | None = None, - objective_thresholds: list[ObjectiveThreshold] | None = None, + objective_thresholds: list[OutcomeConstraint] | None = None, pruning_target_parameterization: Arm | None = None, ) -> None: """Inits MultiObjectiveOptimizationConfig. @@ -323,18 +304,18 @@ def __init__( outcome_constraints=constraints, objective_thresholds=objective_thresholds, ) - self._objective: MultiObjective | ScalarizedObjective = objective + self._objective: Objective = objective self._outcome_constraints: list[OutcomeConstraint] = constraints - self._objective_thresholds: list[ObjectiveThreshold] = objective_thresholds + self._objective_thresholds: list[OutcomeConstraint] = objective_thresholds self.pruning_target_parameterization = pruning_target_parameterization # pyre-fixme[14]: Inconsistent override. def clone_with_args( self, - objective: MultiObjective | ScalarizedObjective | None = None, + objective: Objective | None = None, outcome_constraints: None | (list[OutcomeConstraint]) = _NO_OUTCOME_CONSTRAINTS, objective_thresholds: None - | (list[ObjectiveThreshold]) = _NO_OBJECTIVE_THRESHOLDS, + | (list[OutcomeConstraint]) = _NO_OBJECTIVE_THRESHOLDS, pruning_target_parameterization: Arm | None = _NO_PRUNING_TARGET_PARAMETERIZATION, ) -> "MultiObjectiveOptimizationConfig": @@ -363,12 +344,12 @@ def clone_with_args( ) @property - def objective(self) -> MultiObjective | ScalarizedObjective: + def objective(self) -> Objective: """Get objective.""" return self._objective @objective.setter - def objective(self, objective: MultiObjective | ScalarizedObjective) -> None: + def objective(self, objective: Objective) -> None: """Set objective if not present in outcome constraints.""" self._validate_transformed_optimization_config( objective=objective, @@ -383,13 +364,13 @@ def all_constraints(self) -> list[OutcomeConstraint]: return self.outcome_constraints + self.objective_thresholds @property - def objective_thresholds(self) -> list[ObjectiveThreshold]: + def objective_thresholds(self) -> list[OutcomeConstraint]: """Get objective thresholds.""" return self._objective_thresholds @objective_thresholds.setter def objective_thresholds( - self, objective_thresholds: list[ObjectiveThreshold] + self, objective_thresholds: list[OutcomeConstraint] ) -> None: """Set outcome constraints if valid, else raise.""" self._validate_transformed_optimization_config( @@ -399,17 +380,17 @@ def objective_thresholds( self._objective_thresholds = objective_thresholds @property - def objective_thresholds_dict(self) -> dict[str, ObjectiveThreshold]: + def objective_thresholds_dict(self) -> dict[str, OutcomeConstraint]: """Get a mapping from objective metric name to the corresponding threshold. """ - return {ot.metric.name: ot for ot in self._objective_thresholds} + return {ot.metric_names[0]: ot for ot in self._objective_thresholds} @staticmethod def _validate_transformed_optimization_config( objective: Objective, outcome_constraints: list[OutcomeConstraint] | None = None, - objective_thresholds: list[ObjectiveThreshold] | None = None, + objective_thresholds: list[OutcomeConstraint] | None = None, ) -> None: """Ensure outcome constraints are valid. @@ -424,7 +405,7 @@ def _validate_transformed_optimization_config( outcome_constraints: Constraints to validate. objective_thresholds: Thresholds objectives must exceed. """ - if not isinstance(objective, (MultiObjective, ScalarizedObjective)): + if not (objective.is_multi_objective or objective.is_scalarized_objective): raise TypeError( "`MultiObjectiveOptimizationConfig` requires an objective " "of type `MultiObjective` or `ScalarizedObjective`. " @@ -433,18 +414,24 @@ def _validate_transformed_optimization_config( ) outcome_constraints = outcome_constraints or [] objective_thresholds = objective_thresholds or [] - if isinstance(objective, MultiObjective): - objectives_by_signature = { - obj.metric.signature: obj for obj in objective.objectives - } + if objective.is_multi_objective: + # Build objectives_by_name by decomposing the multi-objective + # expression into sub-objectives + parts = [p.strip() for p in objective.expression.split(",")] + objectives_by_name: dict[str, Objective] = {} + for part in parts: + sub_obj = Objective(expression=part) + for name in sub_obj.metric_names: + objectives_by_name[name] = sub_obj + check_objective_thresholds_match_objectives( - objectives_by_signature=objectives_by_signature, + objectives_by_name=objectives_by_name, objective_thresholds=objective_thresholds, ) - unconstrainable_metrics = objective.get_unconstrainable_metrics() + unconstrainable_metric_names = objective.get_unconstrainable_metric_names() OptimizationConfig._validate_outcome_constraints( - unconstrainable_metrics=unconstrainable_metrics, + unconstrainable_metric_names=unconstrainable_metric_names, outcome_constraints=outcome_constraints, ) @@ -458,29 +445,28 @@ def __repr__(self) -> str: def check_objective_thresholds_match_objectives( - objectives_by_signature: dict[str, Objective], - objective_thresholds: list[ObjectiveThreshold], + objectives_by_name: dict[str, Objective], + objective_thresholds: list[OutcomeConstraint], ) -> None: """Error if thresholds on objective_metrics bound from the wrong direction or if there is a mismatch between objective thresholds and objectives. """ - obj_thresh_metrics = set() + obj_thresh_metrics: set[str] = set() for threshold in objective_thresholds: - th_metric_signature = threshold.metric.signature - th_metric_name = threshold.metric.name - if th_metric_signature not in objectives_by_signature: + th_metric_name = threshold.metric_names[0] + if th_metric_name not in objectives_by_name: raise UserInputError( f"Objective threshold {threshold} is on metric '{th_metric_name}', " f"but that metric is not among the objectives." ) - if th_metric_signature in obj_thresh_metrics: + if th_metric_name in obj_thresh_metrics: raise UserInputError( "More than one objective threshold specified for metric " f"{th_metric_name}." ) - obj_thresh_metrics.add(th_metric_signature) + obj_thresh_metrics.add(th_metric_name) - minimize = objectives_by_signature[th_metric_signature].minimize + minimize = objectives_by_name[th_metric_name].minimize bounded_above = threshold.op == ComparisonOp.LEQ is_aligned = minimize == bounded_above if not is_aligned: @@ -495,7 +481,7 @@ def check_objective_thresholds_match_objectives( class PreferenceOptimizationConfig(MultiObjectiveOptimizationConfig): def __init__( self, - objective: MultiObjective, + objective: Objective, preference_profile_name: str, outcome_constraints: list[OutcomeConstraint] | None = None, expect_relativized_outcomes: bool = False, @@ -529,6 +515,12 @@ def __init__( consideration, and if not, the parameter value will be replaced with the corresponding value in the target arm. """ + if not objective.is_multi_objective: + raise TypeError( + "`PreferenceOptimizationConfig` requires a multi-objective. " + "Use `OptimizationConfig` instead if using a " + "single-metric objective." + ) if outcome_constraints: raise NotImplementedError( "Outcome constraints are not yet supported in " @@ -558,7 +550,7 @@ def is_bope_problem(self) -> bool: # pyre-ignore[14]: Inconsistent override. def clone_with_args( self, - objective: MultiObjective | None = None, + objective: Objective | None = None, preference_profile_name: str | None = None, outcome_constraints: list[OutcomeConstraint] | None = _NO_OUTCOME_CONSTRAINTS, expect_relativized_outcomes: bool | None = None, @@ -566,11 +558,7 @@ def clone_with_args( | None = _NO_PRUNING_TARGET_PARAMETERIZATION, ) -> PreferenceOptimizationConfig: """Make a copy of this optimization config.""" - objective = ( - assert_is_instance(self.objective.clone(), MultiObjective) - if objective is None - else objective - ) + objective = self.objective.clone() if objective is None else objective preference_profile_name = ( self.preference_profile_name diff --git a/ax/core/outcome_constraint.py b/ax/core/outcome_constraint.py index 1d22616f3f4..e35ce3b35bb 100644 --- a/ax/core/outcome_constraint.py +++ b/ax/core/outcome_constraint.py @@ -8,165 +8,165 @@ from __future__ import annotations -import logging -from collections.abc import Iterable +import re +import warnings +from functools import cached_property from ax.core.metric import Metric from ax.core.types import ComparisonOp +from ax.exceptions.core import UserInputError from ax.utils.common.base import SortableBase -from ax.utils.common.logger import get_logger - +from ax.utils.common.string_utils import unsanitize_name +from ax.utils.common.sympy import ( + build_constraint_expression_str, + extract_coefficient_dict_from_inequality, +) -logger: logging.Logger = get_logger(__name__) -CONSTRAINT_WARNING_MESSAGE: str = ( - "Constraint on {name} appears invalid: {bound} bound on metric " - + "for which {is_better} values are better." -) -LOWER_BOUND_MISMATCH: dict[str, str] = {"bound": "Lower", "is_better": "lower"} -UPPER_BOUND_MISMATCH: dict[str, str] = {"bound": "Upper", "is_better": "higher"} +class OutcomeConstraint(SortableBase): + """Class for representing outcome constraints via an expression string. -CONSTRAINT_THRESHOLD_WARNING_MESSAGE: str = ( - "Constraint threshold on {name} appears invalid: {bound} bound on metric " - + "for which {is_better} values are better." -) -UPPER_BOUND_THRESHOLD: dict[str, str] = {"bound": "Positive", "is_better": "lower"} -LOWER_BOUND_THRESHOLD: dict[str, str] = {"bound": "Negative", "is_better": "higher"} + An outcome constraint is an inequality of the form + ``metric >= bound`` or ``metric <= bound``, where the bound can be + absolute or relative. + To indicate a relative constraint (i.e. performance relative to some baseline) + multiply your bound by "baseline". For example "qps >= 0.95 * baseline" will + constrain such that the QPS is at least 95% of the baseline arm's QPS. + The "baseline" in these expressions refers to the experiment's + ``status_quo`` arm. -class OutcomeConstraint(SortableBase): - """Base class for representing outcome constraints. + Examples:: - 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). + OutcomeConstraint(expression="qps >= 700") + OutcomeConstraint(expression="loss <= 0.5") + OutcomeConstraint(expression="2*m1 + 3*m2 <= 10") + OutcomeConstraint(expression="latency <= 0.95 * baseline", ) Attributes: - 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, else ``bound`` 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 in - ``Experiment``. + expression: The expression string defining this constraint. """ def __init__( - self, metric: Metric, op: ComparisonOp, bound: float, relative: bool = True + self, + expression: str | None = None, + *, + relative: bool | None = None, + # Deprecated backward-compat kwargs + metric: Metric | None = None, + op: ComparisonOp | None = None, + bound: float | None = None, ) -> None: - self._validate_metric_constraint_op(metric=metric, op=op) - self._metric = metric - self._op = op - self.bound: float = float(bound) if isinstance(bound, int) else bound - self.relative = relative + """Create a new outcome constraint. + + Args: + expression: A string expression defining the constraint, e.g. + ``"qps >= 700"`` or ``"loss <= 0.5 * baseline"``. + relative: If provided, overrides the relativity inferred from the + expression string. If ``True``, the bound is interpreted as + a percent-difference from the status-quo arm's metric value. + metric: *Deprecated.* A single metric to constrain. Provide + ``expression`` instead. + op: *Deprecated.* Only used with ``metric``. + bound: *Deprecated.* Only used with ``metric``. + """ + if metric is not None: + # Backward-compat path + warnings.warn( + "Passing `metric` to OutcomeConstraint is deprecated. " + "Use `expression` instead, e.g. " + f'OutcomeConstraint(expression="{metric.name} >= 0").', + DeprecationWarning, + stacklevel=2, + ) + if expression is not None: + raise UserInputError("Cannot specify both `expression` and `metric`.") + if op is None: + raise UserInputError( + "Must specify `op` when using the deprecated `metric` kwarg." + ) + if bound is None: + raise UserInputError( + "Must specify `bound` when using the deprecated `metric` kwarg." + ) + inferred_relative = relative if relative is not None else False + expression = build_constraint_expression_str( + metric_weights=[(metric.name, 1.0)], + op=_op_to_str(op), + bound=float(bound), + relative=inferred_relative, + ) + relative = inferred_relative + + if expression is None: + raise UserInputError( + "An expression string is required. " + "Example: OutcomeConstraint(expression='qps >= 700') or " + "OutcomeConstraint(expression='loss <= 0.5')." + ) + + self._expression_str: str = expression + # Eagerly validate the expression so errors surface at construction time. + _ = self._parsed + + @cached_property + def _parsed( + self, + ) -> tuple[list[tuple[str, float]], ComparisonOp, float, bool]: + """Parse the expression string into its constituent parts. + + Returns a tuple of + ``(metric_weights, op, bound, relative)`` where ``metric_weights`` + is a list of ``(metric_name, weight)`` pairs. All public properties + -- ``metric_names``, ``metric_weights``, ``op``, ``bound``, and + ``relative`` -- delegate to this parsed representation. + """ + return _parse_constraint_expression(self._expression_str) + + @property + def expression(self) -> str: + """Get the expression string defining this constraint.""" + return self._expression_str @property - def metric(self) -> Metric: - return self._metric + def metric_names(self) -> list[str]: + """Get a list of all metric names referenced in the constraint.""" + return [name for name, _ in self._parsed[0]] - @metric.setter - def metric(self, metric: Metric) -> None: - self._validate_metric_constraint_op(metric=metric, op=self.op) - self._metric = metric + @property + def metric_weights(self) -> list[tuple[str, float]]: + """Get (metric_name, weight) pairs for this constraint.""" + return list(self._parsed[0]) @property def op(self) -> ComparisonOp: - return self._op + """Get the comparison operator.""" + return self._parsed[1] + + @property + def bound(self) -> float: + """Get the constraint bound.""" + return self._parsed[2] - @op.setter - def op(self, op: ComparisonOp) -> None: - self._validate_metric_constraint_op(metric=self.metric, op=op) - self._op = op + @property + def relative(self) -> bool: + """Whether the bound is relative to the status-quo arm.""" + return self._parsed[3] def clone(self) -> OutcomeConstraint: """Create a copy of this OutcomeConstraint.""" - return OutcomeConstraint( - metric=self.metric.clone(), - op=self.op, - bound=self.bound, - relative=self.relative, - ) + return OutcomeConstraint(expression=self._expression_str) - @staticmethod - def _validate_metric_constraint_op( - metric: Metric, op: ComparisonOp - ) -> tuple[bool, str]: - """Ensure constraint is compatible with metric definition. + def __repr__(self) -> str: + return f"OutcomeConstraint({self._expression_str})" - Args: - metric: Metric to constrain. - op: Specifies whether metric should be greater or equal - to, or less than or equal to, some bound. In case metric has: - - lower_is_better=True, op is interpreted as an upper bound and - should be LEQ. - - lower_is_better=False, op is interpreted as a lower bound and - should be GEQ. - - Returns: A tuple consisting of - - A boolean indicating whether the constraint is valid, - - A string containing a warning message if the constraint is invalid. + def __eq__(self, other: object) -> bool: + if not isinstance(other, OutcomeConstraint): + return False + return self._expression_str == other._expression_str - """ - fmt_data = None - if metric.lower_is_better is not None: - if op == ComparisonOp.GEQ and metric.lower_is_better: - fmt_data = LOWER_BOUND_MISMATCH - if op == ComparisonOp.LEQ and not metric.lower_is_better: - fmt_data = UPPER_BOUND_MISMATCH - if fmt_data is not None: - fmt_data = {**fmt_data, "name": metric.name} - msg = CONSTRAINT_WARNING_MESSAGE.format(**fmt_data) - logger.debug(msg) - return False, msg - return True, "" - - def _validate_constraint(self) -> tuple[bool, str]: - """Ensure constraint is compatible with metric definition. - In case metric has: - - lower_is_better=True: op is interpreted as an upper bound - and should be LEQ; bound should be positive in case - of relative constraint. - - lower_is_better=False, op is interpreted as a lower bound - and should be GEQ; bound should be negative in case - of relative constraint. - - Returns: A tuple consisting of - - A boolean indicating whether the constraint is valid, - - A string containing a warning message if the constraint is invalid. - """ - valid_op, msg = self._validate_metric_constraint_op( - metric=self.metric, op=self.op - ) - if not valid_op: - return False, msg - - if not self.relative: - return True, "" - - fmt_data = None - if self.metric.lower_is_better is not None: - if self.bound < 0 and self.metric.lower_is_better: - fmt_data = UPPER_BOUND_THRESHOLD - if self.bound > 0 and not self.metric.lower_is_better: - fmt_data = LOWER_BOUND_THRESHOLD - if fmt_data is not None: - fmt_data = {**fmt_data, "name": self.metric.name} - msg += CONSTRAINT_THRESHOLD_WARNING_MESSAGE.format(**fmt_data) - logger.debug(msg) - return False, msg - - return True, "" - - def __repr__(self) -> str: - op = ">=" if self.op == ComparisonOp.GEQ else "<=" - relative = "%" if self.relative else "" - return f"OutcomeConstraint({self.metric.name} {op} {self.bound}{relative})" + def __hash__(self) -> int: + return hash(self._expression_str) @property def _unique_id(self) -> str: @@ -174,32 +174,17 @@ def _unique_id(self) -> str: class ObjectiveThreshold(OutcomeConstraint): - """Class for representing Objective Thresholds. + """*Deprecated.* Use ``OutcomeConstraint`` directly. 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). + to contribute to hypervolume calculations. - The direction of the bound is inferred from the Metric's lower_is_better attribute. + Example:: - Attributes: - 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 in - ``Experiment``. - - op: automatically inferred, but manually overwritable. - specifies whether metric should be greater or equal to, or less - than or equal to, some bound. + # Old + ObjectiveThreshold(metric=m, bound=0.5, relative=False) + # New + OutcomeConstraint(expression="m >= 0.5", relative=False) """ def __init__( @@ -209,6 +194,12 @@ def __init__( relative: bool = True, op: ComparisonOp | None = None, ) -> None: + warnings.warn( + "ObjectiveThreshold is deprecated. Use OutcomeConstraint " + "with an expression string instead.", + DeprecationWarning, + stacklevel=2, + ) if metric.lower_is_better is None and op is None: raise ValueError( f"Metric {metric} must have attribute `lower_is_better` set or " @@ -217,49 +208,37 @@ def __init__( elif op is None: op = ComparisonOp.LEQ if metric.lower_is_better else ComparisonOp.GEQ - # It's likely that the metric passed into the ObjectiveThreshold constructor - # is the same instance as the metric in the Objective. Thus, we have to clone - # the metric passed in here to ensure a 1:1 relationship between user-facing - # objects and DB objects. - super().__init__(metric=metric.clone(), op=op, bound=bound, relative=relative) + expression = build_constraint_expression_str( + metric_weights=[(metric.name, 1.0)], + op=_op_to_str(op), + bound=float(bound), + relative=relative, + ) + super().__init__(expression=expression, relative=relative) def clone(self) -> ObjectiveThreshold: """Create a copy of this ObjectiveThreshold.""" - return ObjectiveThreshold( - metric=self.metric.clone(), - bound=self.bound, - relative=self.relative, - op=self.op, - ) + ot = ObjectiveThreshold.__new__(ObjectiveThreshold) + ot._expression_str = self._expression_str + return ot def __repr__(self) -> str: - op = ">=" if self.op == ComparisonOp.GEQ else "<=" - relative = "%" if self.relative else "" - return f"ObjectiveThreshold({self.metric.name} {op} {self.bound}{relative})" + return f"ObjectiveThreshold({self._expression_str})" class ScalarizedOutcomeConstraint(OutcomeConstraint): - """Class for presenting outcome constraints composed of a linear - scalarization of metrics. + """*Deprecated.* Use ``OutcomeConstraint`` with a scalarized expression. - Attributes: - metrics: List of metrics. - weights: Weights for scalarization; default to 1.0 / len(metrics). - 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, else ``bound`` 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 - always 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 in - ``Experiment``. - """ + Example:: - weights: list[float] + # Old + ScalarizedOutcomeConstraint( + metrics=[m1, m2], op=ComparisonOp.LEQ, bound=10, + weights=[2.0, 3.0], + ) + # New + OutcomeConstraint(expression="2.0*m1 + 3.0*m2 <= 10", relative=False) + """ def __init__( self, @@ -269,83 +248,133 @@ def __init__( relative: bool = True, weights: list[float] | None = None, ) -> None: - for metric in metrics: - self._validate_metric_constraint_op(metric=metric, op=op) - + warnings.warn( + "ScalarizedOutcomeConstraint is deprecated. Use OutcomeConstraint " + "with a scalarized expression instead.", + DeprecationWarning, + stacklevel=2, + ) if weights is None: weights = [1.0 / len(metrics)] * len(metrics) elif len(weights) != len(metrics): raise ValueError("Length of weights must equal length of metrics") - self._metrics = metrics - self.weights = weights - self._op = op - self.bound = bound - self.relative = relative - @property - def metric_weights(self) -> Iterable[tuple[Metric, float]]: - """Get the objective metrics and weights.""" - return zip(self.metrics, self.weights) + metric_weights = [(m.name, w) for m, w in zip(metrics, weights)] + expression = build_constraint_expression_str( + metric_weights=metric_weights, + op=_op_to_str(op), + bound=float(bound), + relative=relative, + ) + super().__init__(expression=expression, relative=relative) @property - def metrics(self) -> list[Metric]: - return self._metrics + def weights(self) -> list[float]: + """Get the scalarization weights, derived from the expression.""" + return [w for _, w in self.metric_weights] - @metrics.setter - def metrics(self, metrics: list[Metric]) -> None: - for metric in metrics: - self._validate_metric_constraint_op(metric=metric, op=self.op) - self._metrics = metrics + def clone(self) -> ScalarizedOutcomeConstraint: + """Create a copy of this ScalarizedOutcomeConstraint.""" + cloned = ScalarizedOutcomeConstraint.__new__(ScalarizedOutcomeConstraint) + cloned._expression_str = self._expression_str + return cloned - @property - def metric(self) -> Metric: - """Override base method to error.""" - raise NotImplementedError( - f"{type(self).__name__} is composed of multiple metrics" - ) + def __repr__(self) -> str: + return f"ScalarizedOutcomeConstraint({self._expression_str})" - @metric.setter - def metric(self, metric: Metric) -> None: - """Override base method to error.""" - raise NotImplementedError( - f"{type(self).__name__} is composed of multiple metrics" - ) - @property - def op(self) -> ComparisonOp: - return self._op +def _op_to_str(op: ComparisonOp) -> str: + """Convert a ComparisonOp enum to its string representation.""" + return ">=" if op == ComparisonOp.GEQ else "<=" - @op.setter - def op(self, op: ComparisonOp) -> None: - for metric in self.metrics: - self._validate_metric_constraint_op(metric=metric, op=op) - self._op = op - def clone(self) -> ScalarizedOutcomeConstraint: - """Create a copy of this ScalarizedOutcomeConstraint.""" - return ScalarizedOutcomeConstraint( - metrics=[m.clone() for m in self.metrics], - op=self.op, - bound=self.bound, - relative=self.relative, - weights=self.weights.copy(), - ) +def _parse_constraint_expression( + expression_str: str, +) -> tuple[list[tuple[str, float]], ComparisonOp, float, bool]: + """ + Parse an outcome constraint string into an OutcomeConstraint object using SymPy. + Currently only supports linear constraints of the form "a * x + b * y >= k" or + "a * x + b * y <= k". + + To indicate a relative constraint (i.e. performance relative to some baseline) + multiply your bound by "baseline". For example "qps >= 0.95 * baseline" will + constrain such that the QPS is at least 95% of the baseline arm's QPS. + + Returns: + * A list of (metric_name, weight) pairs. + * The comparison operator. + * The bound. If the constraint is relative, the bound specifies percent- + difference from the observed metric value on the status-quo arm. That is, + the bound's value will be ``(1 + bound / 100.0) * baseline``. + * Whether the bound is relative to the status-quo arm. + """ + # Step 1: Detect and strip "baseline" before sympy parsing. + # This avoids sympy's direction-dependent normalization affecting the + # coefficient sign of the baseline term. + is_relative = False + if "baseline" in expression_str: + is_relative = True + # Try to strip "* baseline" (with optional whitespace) + cleaned = re.sub(r"\s*\*\s*baseline\b", "", expression_str) + if cleaned != expression_str: + expression_str = cleaned + else: + # Standalone "baseline" (no multiplier) means multiplier = 1 + expression_str = re.sub(r"\bbaseline\b", "1", expression_str) + + # Step 2: Parse with sympy (no baseline symbol) + coefficient_dict = extract_coefficient_dict_from_inequality( + inequality_str=expression_str + ) + + # Step 3: Extract metrics and bound + constraint_dict: dict[str, float] = {} + bound = 0.0 + for term, coefficient in coefficient_dict.items(): + if term.is_symbol: + constraint_dict[term.name] = coefficient + elif term.is_number: + # Invert because we are "moving" the bound to the right hand side + bound = -1 * coefficient + else: + raise UserInputError( + f"Only linear outcome constraints are supported, found {expression_str}" + ) + + # Step 4: Normalize and convert multiplier -> percentage for relative + if len(constraint_dict) == 1: + term, coefficient = next(iter(constraint_dict.items())) + bound = bound / coefficient + 0.0 # +0.0 normalizes -0.0 to 0.0 + if is_relative: + bound = round((bound - 1) * 100, 10) - def __repr__(self) -> str: - terms = [] - for metric, weight in zip(self.metrics, self.weights): - # For the first term, keep the sign only if it's negative - if len(terms) == 0: - terms.append(f"{weight} * {metric.name}") - else: - if weight < 0: - terms.append(f"- {abs(weight)} * {metric.name}") - else: - terms.append(f"+ {weight} * {metric.name}") - - op = ">=" if self.op == ComparisonOp.GEQ else "<=" - relative = "%" if self.relative else "" return ( - "ScalarizedOutcomeConstraint(" - f"{' '.join(terms)} {op} {self.bound}{relative})" + [(unsanitize_name(term), 1)], + ComparisonOp.LEQ if coefficient > 0 else ComparisonOp.GEQ, + bound, + is_relative, ) + + names, coefficients = zip(*constraint_dict.items()) + + # Detect GEQ: sympy normalizes all inequalities to LEQ by negating both + # sides for GEQ. If all coefficients are negative, flip back. + if all(c < 0 for c in coefficients): + op = ComparisonOp.GEQ + coefficients = tuple(-c for c in coefficients) + bound = -bound + else: + op = ComparisonOp.LEQ + + if is_relative: + bound = round((bound - 1) * 100, 10) + + return ( + [ + (unsanitize_name(name), coefficient) + for name, coefficient in zip(names, coefficients) + ], + op, + bound, + is_relative, + ) diff --git a/ax/core/tests/test_derived_metric.py b/ax/core/tests/test_derived_metric.py index 9618cc34e96..e3660206f54 100644 --- a/ax/core/tests/test_derived_metric.py +++ b/ax/core/tests/test_derived_metric.py @@ -470,7 +470,7 @@ def test_two_phase_experiment_fetch(self) -> None: ) ], ), - tracking_metrics=[Metric(name="base_a"), Metric(name="base_b")], + tracking_metrics=[Metric(name="base_a"), Metric(name="base_b"), derived], ) self.assertIsInstance(experiment.metrics["total"], DerivedMetric) @@ -923,7 +923,11 @@ def test_experiment_integration(self) -> None: ) ], ), - tracking_metrics=[Metric(name="base_a"), Metric(name="base_b")], + tracking_metrics=[ + Metric(name="base_a"), + Metric(name="base_b"), + derived_ratio, + ], ) for i in range(3): diff --git a/ax/core/tests/test_experiment.py b/ax/core/tests/test_experiment.py index 9b7aaca9fda..ce9ca6fd5dd 100644 --- a/ax/core/tests/test_experiment.py +++ b/ax/core/tests/test_experiment.py @@ -52,6 +52,7 @@ from ax.metrics.branin import BraninMetric from ax.metrics.hartmann6 import Hartmann6Metric from ax.metrics.noisy_function import NoisyFunctionMetric +from ax.metrics.noisy_function_map import NoisyFunctionMapMetric from ax.runners.synthetic import SyntheticRunner from ax.service.ax_client import AxClient from ax.service.utils.instantiation import ObjectiveProperties @@ -190,7 +191,8 @@ def test_tracking_metrics_merge(self) -> None: tracking_metrics=[Metric(name="m1"), Metric(name="m3")], ) self.assertEqual( - len(none_throws(exp.optimization_config).metrics) + 1, len(exp.metrics) + len(none_throws(exp.optimization_config).metric_names) + 1, + len(exp.metrics), ) def test_basic_batch_creation(self) -> None: @@ -293,31 +295,43 @@ def test_only_range_parameter_constraints(self) -> None: def test_metric_setters(self) -> None: # Establish current metrics size + # get_optimization_config() has metric_names {"m1", "m2"}, len=2 + # experiment has m1, m2, tracking = 3 self.assertEqual( - len(get_optimization_config().metrics) + 1, len(self.experiment.metrics) + len(get_optimization_config().metric_names) + 1, + len(self.experiment.metrics), ) - # Add optimization config with 1 different metric - opt_config = get_optimization_config() - opt_config.outcome_constraints[0].metric = Metric(name="m3") + # Add a new metric and set optimization config using it as constraint + self.experiment.add_metric(Metric(name="m3")) + opt_config = OptimizationConfig( + objective=Objective(expression="m1"), + outcome_constraints=[ + OutcomeConstraint(expression="m3 >= -0.25 * baseline") + ], + ) self.experiment.optimization_config = opt_config - # Verify total metrics has increaed by 1. + # Verify total metrics has increased by 1 (m1, m2, m3, tracking = 4). self.assertEqual( - len(get_optimization_config().metrics) + 2, len(self.experiment.metrics) + len(get_optimization_config().metric_names) + 2, + len(self.experiment.metrics), ) # Add optimization config with 1 scalarized constraint composed of 2 metrics + self.experiment.add_metric(Metric(name="oc_m3")) + self.experiment.add_metric(Metric(name="oc_m4")) opt_config = get_optimization_config() opt_config.outcome_constraints = opt_config.outcome_constraints + [ get_scalarized_outcome_constraint() ] self.experiment.optimization_config = opt_config - # Verify total metrics size is the same. - self.assertEqual(len(opt_config.metrics) + 2, len(self.experiment.metrics)) + # Verify total metrics size (m1, m2, m3, oc_m3, oc_m4, tracking = 6). + self.assertEqual(len(opt_config.metric_names) + 2, len(self.experiment.metrics)) self.assertEqual( - len(get_optimization_config().metrics) + 4, len(self.experiment.metrics) + len(get_optimization_config().metric_names) + 4, + len(self.experiment.metrics), ) # set back self.experiment.optimization_config = get_optimization_config() @@ -325,13 +339,15 @@ def test_metric_setters(self) -> None: # Test adding new tracking metric self.experiment.add_tracking_metric(Metric(name="m4")) self.assertEqual( - len(get_optimization_config().metrics) + 5, len(self.experiment.metrics) + len(get_optimization_config().metric_names) + 5, + len(self.experiment.metrics), ) # Test adding new tracking metrics self.experiment.add_tracking_metrics([Metric(name="z1")]) self.assertEqual( - len(get_optimization_config().metrics) + 6, len(self.experiment.metrics) + len(get_optimization_config().metric_names) + 6, + len(self.experiment.metrics), ) # Verify update_tracking_metric updates the metric definition @@ -347,11 +363,11 @@ def test_metric_setters(self) -> None: with self.assertRaises(ValueError): self.experiment.add_tracking_metrics([Metric(name="z1"), Metric(name="m4")]) - # Verify unable to add metric in optimization config + # Verify unable to add metric already on experiment with self.assertRaises(ValueError): self.experiment.add_tracking_metric(Metric(name="m1")) - # Verify unable to add metric in optimization config + # Verify unable to add metric already on experiment with self.assertRaises(ValueError): self.experiment.add_tracking_metrics([Metric(name="z2"), Metric(name="m1")]) @@ -504,13 +520,21 @@ def test_disable_search_space_parameters(self) -> None: def test_optimization_config_setter(self) -> None: # Establish current metrics size self.assertEqual( - len(get_optimization_config().metrics) + 1, len(self.experiment.metrics) + len(get_optimization_config().metric_names) + 1, + len(self.experiment.metrics), ) - # Add optimization config with 1 different metric + # Setting an opt config whose metrics are all registered should work opt_config = get_optimization_config() - opt_config.outcome_constraints[0].metric = Metric(name="m3") - self + self.experiment.optimization_config = opt_config + + # Setting an opt config with an unregistered metric should raise + new_opt_config = OptimizationConfig( + objective=Objective(expression="m1"), + outcome_constraints=[OutcomeConstraint(expression="unknown_metric >= 0.5")], + ) + with self.assertRaisesRegex(ValueError, "not found on experiment"): + self.experiment.optimization_config = new_opt_config def test_status_quo_setter(self) -> None: sq_parameters = self.experiment.status_quo.parameters @@ -756,10 +780,6 @@ def test_bulk_configure_metrics(self) -> None: attributes_to_update={"param_names": ["foo", "bar"]}, ) self.assertEqual(inital_metrics, exp.metrics) - # double check the change is populated in opt config too - inital_metrics.pop("no update") - self.assertIsNotNone(exp.optimization_config) - self.assertEqual(inital_metrics, exp.optimization_config.metrics) with self.subTest("parameter not in initialization"): with self.assertRaisesRegex( UserInputError, "does not contain the requested " @@ -1264,7 +1284,12 @@ def test_warm_start_from_old_experiment(self) -> None: else: trial.mark_completed() # make metric noiseless for exact reproducibility - old_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + old_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + old_experiment.get_metric(_obj_name), NoisyFunctionMetric + ).noise_sd = 0 old_experiment.fetch_data() # should fail if new_experiment has trials @@ -1281,7 +1306,12 @@ def test_warm_start_from_old_experiment(self) -> None: # check that all non-failed trials are copied to new_experiment new_experiment = get_branin_experiment() # make metric noiseless for exact reproducibility - new_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + new_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + new_experiment.get_metric(_obj_name), NoisyFunctionMetric + ).noise_sd = 0 for _, trial in old_experiment.trials.items(): trial._run_metadata = DUMMY_RUN_METADATA # name one arm to test name-preserving logic. @@ -1327,7 +1357,12 @@ def test_warm_start_from_old_experiment(self) -> None: # check that all non-failed/abandoned trials are copied to new_experiment new_experiment = get_branin_experiment() # make metric noiseless for exact reproducibility - new_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + new_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + new_experiment.get_metric(_obj_name), NoisyFunctionMetric + ).noise_sd = 0 new_experiment.warm_start_from_old_experiment( old_experiment=old_experiment, trial_statuses_to_copy=[TrialStatus.COMPLETED], @@ -1339,7 +1374,12 @@ def test_warm_start_from_old_experiment(self) -> None: # check that only run_metadata of specified keys are copied new_experiment = get_branin_experiment() - new_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + new_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + new_experiment.get_metric(_obj_name), NoisyFunctionMetric + ).noise_sd = 0 new_experiment.warm_start_from_old_experiment( old_experiment=old_experiment, copy_run_metadata_keys=[DUMMY_RUN_METADATA_KEY_1], @@ -1526,6 +1566,12 @@ def test_metric_summary_df(self) -> None: ], ), tracking_metrics=[ + # Pass opt config metrics with their real types so the + # experiment preserves the Metric subclass information. + Metric(name="my_objective_1", lower_is_better=True), + TestMetric(name="my_objective_2"), + Metric(name="my_constraint_1", lower_is_better=False), + TestMetric(name="my_constraint_2"), Metric(name="my_tracking_metric_1", lower_is_better=True), TestMetric(name="my_tracking_metric_2", lower_is_better=False), Metric(name="my_tracking_metric_3"), @@ -1661,6 +1707,7 @@ def test_trial_indices_with_data(self) -> None: trials = exp.trial_indices_with_data(critical_metrics_only=True) self.assertEqual(trials, {0, 1}) with self.subTest("changed opt config, no data for new config"): + exp.add_metric(Metric(name="test_metric")) exp.optimization_config = get_optimization_config_no_constraints() trials = exp.trial_indices_with_data(critical_metrics_only=True) self.assertEqual(len(trials), 0) @@ -2207,13 +2254,23 @@ def test_warm_start_map_data(self) -> None: else: trial.mark_completed() # make metric noiseless for exact reproducibility - old_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + old_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + old_experiment.get_metric(_obj_name), NoisyFunctionMapMetric + ).noise_sd = 0 old_experiment.fetch_data() # check that all non-failed trials are copied to new_experiment new_experiment = get_branin_experiment_with_timestamp_map_metric() # make metric noiseless for exact reproducibility - new_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + new_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + new_experiment.get_metric(_obj_name), NoisyFunctionMapMetric + ).noise_sd = 0 for _, trial in old_experiment.trials.items(): trial._run_metadata = DUMMY_RUN_METADATA new_experiment.warm_start_from_old_experiment(old_experiment=old_experiment) @@ -2239,7 +2296,12 @@ def test_warm_start_map_data(self) -> None: # check that only run_metadata of specified keys are copied new_experiment = get_branin_experiment_with_timestamp_map_metric() - new_experiment.optimization_config.objective.metric.noise_sd = 0 + _obj_name = none_throws( + new_experiment.optimization_config + ).objective.metric_names[0] + assert_is_instance( + new_experiment.get_metric(_obj_name), NoisyFunctionMapMetric + ).noise_sd = 0 for _, trial in old_experiment.trials.items(): trial._run_metadata = DUMMY_RUN_METADATA new_experiment.warm_start_from_old_experiment( @@ -2498,6 +2560,8 @@ def test_is_bope_problem(self) -> None: with self.subTest("Has PreferenceOptimizationConfig"): experiment = get_branin_experiment() + experiment.add_metric(Metric(name="m1")) + experiment.add_metric(Metric(name="m2")) pref_opt_config = PreferenceOptimizationConfig( objective=MultiObjective( objectives=[ diff --git a/ax/core/tests/test_multi_type_experiment.py b/ax/core/tests/test_multi_type_experiment.py index d8a28e60ac2..b314cfd5b75 100644 --- a/ax/core/tests/test_multi_type_experiment.py +++ b/ax/core/tests/test_multi_type_experiment.py @@ -151,8 +151,10 @@ def test_setting_opt_config(self) -> None: self.assertDictEqual( self.experiment._metric_to_trial_type, {"m1": "type1", "m2": "type2"} ) + m3 = BraninMetric("m3", ["x1", "x2"]) + self.experiment.add_tracking_metric(m3) self.experiment.optimization_config = OptimizationConfig( - Objective(BraninMetric("m3", ["x1", "x2"]), minimize=True) + Objective(metric=m3, minimize=True) ) self.assertDictEqual( self.experiment._metric_to_trial_type, diff --git a/ax/core/tests/test_objective.py b/ax/core/tests/test_objective.py index 6e28cc39b09..de434f2c22a 100644 --- a/ax/core/tests/test_objective.py +++ b/ax/core/tests/test_objective.py @@ -6,6 +6,8 @@ # pyre-strict +import warnings + from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective, ScalarizedObjective from ax.exceptions.core import UserInputError @@ -20,76 +22,196 @@ def setUp(self) -> None: "m2": Metric(name="m2", lower_is_better=True), "m3": Metric(name="m3", lower_is_better=False), } - self.objectives = { - "o1": Objective(metric=self.metrics["m1"], minimize=True), - "o2": Objective(metric=self.metrics["m2"], minimize=True), - "o3": Objective(metric=self.metrics["m3"], minimize=False), - } - self.objective = Objective(metric=self.metrics["m1"], minimize=False) - self.multi_objective = MultiObjective( - objectives=[ - self.objectives["o1"], - self.objectives["o2"], - self.objectives["o3"], - ] - ) - self.scalarized_objective = ScalarizedObjective( - metrics=[self.metrics["m1"], self.metrics["m2"]], minimize=True - ) + # Expression-based objectives + self.objective = Objective(expression="m1") + self.objective_minimize = Objective(expression="-m1") + self.multi_objective = Objective(expression="-m1, -m2, m3") + self.scalarized_objective = Objective(expression="2*m1 + m2") + + def test_ExpressionInit(self) -> None: + """Test creating objectives with expression strings.""" + # Single metric maximize + obj = Objective(expression="accuracy") + self.assertEqual(obj.expression, "accuracy") + self.assertEqual(obj.metric_names, ["accuracy"]) + self.assertFalse(obj.minimize) + self.assertTrue(obj.is_single_objective) + self.assertFalse(obj.is_multi_objective) + self.assertFalse(obj.is_scalarized_objective) + + # Single metric minimize + obj_min = Objective(expression="-loss") + self.assertEqual(obj_min.expression, "-loss") + self.assertEqual(obj_min.metric_names, ["loss"]) + self.assertTrue(obj_min.minimize) + + # Scalarized + obj_scalar = Objective(expression="2*acc + recall") + self.assertEqual(obj_scalar.metric_names, ["acc", "recall"]) + self.assertTrue(obj_scalar.is_scalarized_objective) + self.assertFalse(obj_scalar.is_multi_objective) + self.assertFalse(obj_scalar.is_single_objective) + self.assertEqual(obj_scalar.metric_weights, [("acc", 2.0), ("recall", 1.0)]) + + # Multi-objective + obj_multi = Objective(expression="acc, -loss") + self.assertEqual(obj_multi.metric_names, ["acc", "loss"]) + self.assertTrue(obj_multi.is_multi_objective) + self.assertFalse(obj_multi.is_single_objective) + self.assertFalse(obj_multi.is_scalarized_objective) + + # No expression should error + with self.assertRaisesRegex(UserInputError, "expression string is required"): + Objective() + with self.assertRaisesRegex(UserInputError, "expression string is required"): + Objective(expression=None) + + # Cannot specify both expression and metric + with self.assertRaisesRegex(UserInputError, "Cannot specify both"): + Objective(expression="m1", metric=self.metrics["m1"]) + + def test_DeprecatedMetricInit(self) -> None: + """Test the deprecated metric-based init path.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + obj = Objective(metric=self.metrics["m2"], minimize=True) + self.assertEqual(len(w), 1) + self.assertIn("deprecated", str(w[0].message).lower()) + self.assertEqual(obj.expression, "-m2") + self.assertTrue(obj.minimize) + self.assertEqual(obj.metric_names, ["m2"]) - def test_Init(self) -> None: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + obj2 = Objective(metric=self.metrics["m3"], minimize=False) + self.assertEqual(len(w), 1) + self.assertEqual(obj2.expression, "m3") + self.assertFalse(obj2.minimize) + + # Test lower_is_better validation still works with self.assertRaisesRegex(UserInputError, "does not specify"): Objective(metric=self.metrics["m1"]) + with self.assertRaisesRegex( UserInputError, "doesn't match the specified optimization direction" ): Objective(metric=self.metrics["m2"], minimize=False) - with self.assertRaises(ValueError): - ScalarizedObjective( - metrics=[self.metrics["m1"], self.metrics["m2"]], weights=[1.0] - ) - with self.assertRaisesRegex( - ValueError, - "Metric with name m2 specifies `lower_is_better` = " - "True, which doesn't match the specified optimization direction.", - ): - # Should fail since m2 specifies lower_is_better=True - ScalarizedObjective( - metrics=[self.metrics["m1"], self.metrics["m2"]], - minimize=False, - ) + + def test_MetricWeights(self) -> None: + """Test metric_weights property.""" + obj = Objective(expression="m1") + self.assertEqual(obj.metric_weights, [("m1", 1.0)]) + + obj_neg = Objective(expression="-m1") + self.assertEqual(obj_neg.metric_weights, [("m1", -1.0)]) + + obj_scalar = Objective(expression="2*m1 + 3*m2") + self.assertEqual(obj_scalar.metric_weights, [("m1", 2.0), ("m2", 3.0)]) + + # metric_weights works for multi-objective too (flat list) + obj_multi = Objective(expression="m1, m2") + self.assertEqual(obj_multi.metric_weights, [("m1", 1.0), ("m2", 1.0)]) + + obj_multi2 = Objective(expression="2*m1, -m2") + self.assertEqual(obj_multi2.metric_weights, [("m1", 2.0), ("m2", -1.0)]) + + def test_Minimize(self) -> None: + """Test minimize property.""" + self.assertFalse(Objective(expression="m1").minimize) + self.assertTrue(Objective(expression="-m1").minimize) + + # Should raise for scalarized + with self.assertRaisesRegex(UserInputError, "single-metric"): + Objective(expression="m1 + m2").minimize + + # Should raise for multi + with self.assertRaisesRegex(UserInputError, "single-metric"): + Objective(expression="m1, m2").minimize + + def test_Clone(self) -> None: + """Test clone method.""" + obj = Objective(expression="2*m1 + m2") + cloned = obj.clone() + self.assertEqual(obj, cloned) + self.assertIsNot(obj, cloned) + self.assertEqual(obj.expression, cloned.expression) + + def test_Repr(self) -> None: + """Test __repr__.""" + self.assertEqual(str(self.objective), 'Objective(expression="m1")') + self.assertEqual(str(self.objective_minimize), 'Objective(expression="-m1")') + + def test_GetUnconstrainableMetricNames(self) -> None: + """Test get_unconstrainable_metric_names.""" + self.assertEqual(self.objective.get_unconstrainable_metric_names(), ["m1"]) self.assertEqual( - self.objective.get_unconstrainable_metrics(), [self.metrics["m1"]] + self.multi_objective.get_unconstrainable_metric_names(), + ["m1", "m2", "m3"], ) + def test_DeprecatedProperties(self) -> None: + """Test that removed deprecated properties raise AttributeError.""" + + # Use a helper to prevent Pyre from resolving the attribute name + # statically; these attributes were intentionally removed. + def _getattr(obj: object, name: str) -> object: + return getattr(obj, name) + + with self.assertRaises(AttributeError): + _getattr(self.objective, "metric") + + with self.assertRaises(AttributeError): + _getattr(self.objective, "metrics") + + with self.assertRaises(AttributeError): + _getattr(self.objective, "metric_signatures") + + with self.assertRaises(AttributeError): + _getattr(self.objective, "get_unconstrainable_metrics") + def test_MultiObjective(self) -> None: - with self.assertRaises(NotImplementedError): - self.multi_objective.metric + """Test deprecated MultiObjective subclass.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + objectives = [ + Objective(expression="-m1"), + Objective(expression="-m2"), + Objective(expression="m3"), + ] + mo = MultiObjective(objectives=objectives) + # Check deprecation warning was emitted + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + self.assertTrue(len(deprecation_warnings) > 0) - self.assertEqual(self.multi_objective.metrics, list(self.metrics.values())) - minimizes = [obj.minimize for obj in self.multi_objective.objectives] - self.assertEqual(minimizes, [True, True, False]) - self.assertEqual(self.multi_objective.clone(), self.multi_objective) - self.assertEqual( - str(self.multi_objective), - ( - "MultiObjective(objectives=" - '[Objective(metric_name="m1", minimize=True), ' - 'Objective(metric_name="m2", minimize=True), ' - 'Objective(metric_name="m3", minimize=False)])' - ), - ) - self.assertEqual( - self.multi_objective.get_unconstrainable_metrics(), - [self.metrics["m1"], self.metrics["m2"], self.metrics["m3"]], - ) + self.assertTrue(mo.is_multi_objective) + self.assertEqual(mo.metric_names, ["m1", "m2", "m3"]) + + # Verify metric_weights reflect the sub-objective directions + self.assertEqual(mo.metric_weights, [("m1", -1.0), ("m2", -1.0), ("m3", 1.0)]) + + # Clone returns a plain Objective with the same expression + cloned = mo.clone() + self.assertEqual(mo.expression, cloned.expression) + self.assertEqual(mo.metric_names, cloned.metric_names) + + # Test repr + self.assertEqual(str(mo), 'Objective(expression="-m1, -m2, m3")') # Test that ScalarizedObjective is not allowed in MultiObjective + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + scalarized = ScalarizedObjective( + metrics=[self.metrics["m1"], self.metrics["m2"]], + minimize=True, + ) + with self.assertRaisesRegex( NotImplementedError, "Scalarized objectives are not supported for a `MultiObjective`.", ): - MultiObjective(objectives=[self.scalarized_objective]) + MultiObjective(objectives=[scalarized]) with self.assertRaisesRegex( NotImplementedError, @@ -97,70 +219,121 @@ def test_MultiObjective(self) -> None: ): MultiObjective( objectives=[ - self.objectives["o1"], - self.scalarized_objective, - self.objectives["o3"], + Objective(expression="-m1"), + scalarized, + Objective(expression="m3"), ] ) def test_ScalarizedObjective(self) -> None: - with self.assertRaises(NotImplementedError): - self.scalarized_objective.metric + """Test deprecated ScalarizedObjective subclass.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + so = ScalarizedObjective( + metrics=[self.metrics["m1"], self.metrics["m2"]], + minimize=True, + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + self.assertTrue(len(deprecation_warnings) > 0) - self.assertEqual( - self.scalarized_objective.metrics, [self.metrics["m1"], self.metrics["m2"]] - ) - weights = [mw[1] for mw in self.scalarized_objective.metric_weights] - self.assertEqual(weights, [1.0, 1.0]) - self.assertEqual(self.scalarized_objective.clone(), self.scalarized_objective) - self.assertEqual( - str(self.scalarized_objective), - ( - "ScalarizedObjective(metric_names=['m1', 'm2'], weights=[1.0, 1.0], " - "minimize=True)" - ), - ) - self.assertEqual( - self.scalarized_objective.get_unconstrainable_metrics(), - [self.metrics["m1"], self.metrics["m2"]], - ) + # Should build a negated expression since minimize=True + self.assertEqual(so.expression, "-m1 - m2") + self.assertEqual(so.metric_names, ["m1", "m2"]) + + # metric_weights should return (name, weight) tuples + self.assertEqual(so.metric_weights, [("m1", -1.0), ("m2", -1.0)]) + + # Clone returns a plain Objective with the same expression + cloned = so.clone() + self.assertEqual(so.expression, cloned.expression) + self.assertEqual(so.metric_names, cloned.metric_names) + + # Weight validation + with self.assertRaises(ValueError): + ScalarizedObjective( + metrics=[self.metrics["m1"], self.metrics["m2"]], weights=[1.0] + ) + + # lower_is_better validation + with self.assertRaisesRegex( + ValueError, + "Metric with name m2 specifies `lower_is_better` = " + "True, which doesn't match the specified optimization direction.", + ): + ScalarizedObjective( + metrics=[self.metrics["m1"], self.metrics["m2"]], + minimize=False, + ) def test_ScalarizedObjective_expression(self) -> None: """Test the expression property of ScalarizedObjective.""" test_cases = [ - # (name, metrics, weights, expected_expression) + # (name, metrics, weights, minimize, expected_expression) ( - "default_weights", + "default_weights_maximize", [self.metrics["m1"], self.metrics["m3"]], None, + False, "m1 + m3", ), ( - "custom_weights", + "custom_weights_maximize", [self.metrics["m1"], self.metrics["m3"]], [1.0, 2.0], + False, "m1 + 2.0*m3", ), ( "negative_weight", [Metric(name="accuracy"), Metric(name="latency")], [1.0, -0.5], + False, "accuracy - 0.5*latency", ), ( "negative_one_weight", [Metric(name="reward"), Metric(name="cost")], [2.0, -1.0], + False, "2.0*reward - cost", ), ] - for name, metrics, weights, expected in test_cases: + for name, metrics, weights, minimize, expected in test_cases: with self.subTest(name=name, expected=expected): - if weights is None: - obj = ScalarizedObjective(metrics=metrics, minimize=False) - else: - obj = ScalarizedObjective( - metrics=metrics, weights=weights, minimize=False - ) + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + if weights is None: + obj = ScalarizedObjective(metrics=metrics, minimize=minimize) + else: + obj = ScalarizedObjective( + metrics=metrics, weights=weights, minimize=minimize + ) self.assertEqual(obj.expression, expected) + + def test_IsInstanceChecks(self) -> None: + """Test that isinstance checks work for deprecated subclasses.""" + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + mo = MultiObjective( + objectives=[ + Objective(expression="m1"), + Objective(expression="-m2"), + ] + ) + so = ScalarizedObjective( + metrics=[self.metrics["m1"], self.metrics["m3"]], + minimize=False, + ) + + self.assertIsInstance(mo, Objective) + self.assertIsInstance(mo, MultiObjective) + self.assertIsInstance(so, Objective) + self.assertIsInstance(so, ScalarizedObjective) + + def test_UniqueId(self) -> None: + """Test _unique_id used for sorting.""" + obj = Objective(expression="m1") + self.assertEqual(obj._unique_id, str(obj)) diff --git a/ax/core/tests/test_optimization_config.py b/ax/core/tests/test_optimization_config.py index b2ce7cefd18..29f253998d0 100644 --- a/ax/core/tests/test_optimization_config.py +++ b/ax/core/tests/test_optimization_config.py @@ -6,6 +6,8 @@ # pyre-strict +import warnings + from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective, ScalarizedObjective from ax.core.optimization_config import ( @@ -25,18 +27,18 @@ OC_STR = ( - 'OptimizationConfig(objective=Objective(metric_name="m1", minimize=False), ' - "outcome_constraints=[OutcomeConstraint(m3 >= -0.25%), " - "OutcomeConstraint(m4 <= 0.25%), " - "ScalarizedOutcomeConstraint(0.5 * m3 + 0.5 * m4 >= -0.25%)])" + "OptimizationConfig(" + 'objective=Objective(expression="m1"), ' + "outcome_constraints=[OutcomeConstraint(m3 >= -0.25), " + "OutcomeConstraint(m4 <= 0.25), " + "ScalarizedOutcomeConstraint(0.5*m3 + 0.5*m4 >= 0.9975 * baseline)])" ) MOOC_STR = ( - "MultiObjectiveOptimizationConfig(objective=MultiObjective(objectives=" - '[Objective(metric_name="m1", minimize=True), ' - 'Objective(metric_name="m2", minimize=False)]), ' - "outcome_constraints=[OutcomeConstraint(m3 >= -0.25%), " - "OutcomeConstraint(m3 <= 0.25%)], objective_thresholds=[])" + "MultiObjectiveOptimizationConfig(" + 'objective=Objective(expression="-m1, m2"), ' + "outcome_constraints=[OutcomeConstraint(m3 >= -0.25), " + "OutcomeConstraint(m3 <= 0.25)], objective_thresholds=[])" ) @@ -49,26 +51,28 @@ def setUp(self) -> None: "m3": Metric(name="m3"), "m4": Metric(name="m4"), } - self.objective = Objective(metric=self.metrics["m1"], minimize=False) - self.alt_objective = Objective(metric=self.metrics["m3"], minimize=False) - self.multi_objective = MultiObjective( - objectives=[self.objective, self.alt_objective], - ) - self.m2_objective = ScalarizedObjective( - metrics=[self.metrics["m1"], self.metrics["m2"]] - ) - self.outcome_constraint = OutcomeConstraint( - metric=self.metrics["m3"], op=ComparisonOp.GEQ, bound=-0.25 - ) - self.additional_outcome_constraint = OutcomeConstraint( - metric=self.metrics["m4"], op=ComparisonOp.LEQ, bound=0.25 - ) - self.scalarized_outcome_constraint = ScalarizedOutcomeConstraint( - metrics=[self.metrics["m3"], self.metrics["m4"]], - weights=[0.5, 0.5], - op=ComparisonOp.GEQ, - bound=-0.25, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + self.objective = Objective(metric=self.metrics["m1"], minimize=False) + self.alt_objective = Objective(metric=self.metrics["m3"], minimize=False) + self.multi_objective = MultiObjective( + objectives=[self.objective, self.alt_objective], + ) + self.m2_objective = ScalarizedObjective( + metrics=[self.metrics["m1"], self.metrics["m2"]] + ) + self.outcome_constraint = OutcomeConstraint( + metric=self.metrics["m3"], op=ComparisonOp.GEQ, bound=-0.25 + ) + self.additional_outcome_constraint = OutcomeConstraint( + metric=self.metrics["m4"], op=ComparisonOp.LEQ, bound=0.25 + ) + self.scalarized_outcome_constraint = ScalarizedOutcomeConstraint( + metrics=[self.metrics["m3"], self.metrics["m4"]], + weights=[0.5, 0.5], + op=ComparisonOp.GEQ, + bound=-0.25, + ) self.outcome_constraints = [ self.outcome_constraint, self.additional_outcome_constraint, @@ -84,7 +88,7 @@ def test_Init(self) -> None: config1.objective = self.alt_objective # constrained Objective. # updating constraints is fine. config1.outcome_constraints = [self.outcome_constraint] - self.assertEqual(len(config1.metrics), 2) + self.assertEqual(len(config1.metric_names), 2) # objective without outcome_constraints is also supported config2 = OptimizationConfig(objective=self.objective) @@ -108,9 +112,11 @@ def test_Eq(self) -> None: ) self.assertEqual(config1, config2) - new_outcome_constraint = OutcomeConstraint( - metric=self.metrics["m2"], op=ComparisonOp.LEQ, bound=0.5 - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + new_outcome_constraint = OutcomeConstraint( + metric=self.metrics["m2"], op=ComparisonOp.LEQ, bound=0.5 + ) config3 = OptimizationConfig( objective=self.objective, outcome_constraints=[self.outcome_constraint, new_outcome_constraint], @@ -123,9 +129,11 @@ def test_ConstraintValidation(self) -> None: OptimizationConfig(objective=self.multi_objective) # Can't constrain on objective metric. - objective_constraint = OutcomeConstraint( - metric=self.objective.metric, op=ComparisonOp.GEQ, bound=0 - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + objective_constraint = OutcomeConstraint( + metric=self.metrics["m1"], op=ComparisonOp.GEQ, bound=0 + ) with self.assertRaises(ValueError): OptimizationConfig( objective=self.objective, outcome_constraints=[objective_constraint] @@ -140,11 +148,13 @@ def test_ConstraintValidation(self) -> None: ) # Two outcome_constraints on the same metric with the same op # should raise. - duplicate_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - op=self.outcome_constraint.op, - bound=self.outcome_constraint.bound + 1, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + duplicate_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + op=self.outcome_constraint.op, + bound=self.outcome_constraint.bound + 1, + ) with self.assertRaises(ValueError): OptimizationConfig( objective=self.objective, @@ -152,12 +162,14 @@ def test_ConstraintValidation(self) -> None: ) # Three outcome_constraints on the same metric should raise. - opposing_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. - op=not self.outcome_constraint.op, - bound=self.outcome_constraint.bound, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + opposing_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. + op=not self.outcome_constraint.op, + bound=self.outcome_constraint.bound, + ) with self.assertRaises(ValueError): OptimizationConfig( objective=self.objective, @@ -167,12 +179,14 @@ def test_ConstraintValidation(self) -> None: # Two outcome_constraints on the same metric with different ops and # flipped bounds (lower < upper) should raise. add_bound = 1 if self.outcome_constraint.op == ComparisonOp.LEQ else -1 - opposing_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. - op=not self.outcome_constraint.op, - bound=self.outcome_constraint.bound + add_bound, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + opposing_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. + op=not self.outcome_constraint.op, + bound=self.outcome_constraint.bound + add_bound, + ) with self.assertRaises(ValueError): OptimizationConfig( objective=self.objective, @@ -181,12 +195,14 @@ def test_ConstraintValidation(self) -> None: # Two outcome_constraints on the same metric with different ops and # bounds should not raise. - opposing_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. - op=not self.outcome_constraint.op, - bound=self.outcome_constraint.bound + 1, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + opposing_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. + op=not self.outcome_constraint.op, + bound=self.outcome_constraint.bound + 1, + ) config = OptimizationConfig( objective=self.objective, outcome_constraints=([self.outcome_constraint, opposing_constraint]), @@ -205,12 +221,17 @@ def test_ConstraintValidation(self) -> None: # Can't constrain on metric in ScalarizedOutcomeConstraint # that overlaps with objective - scalarized_with_objective_metric = ScalarizedOutcomeConstraint( - metrics=[self.metrics["m1"], self.metrics["m4"]], # m1 is objective metric - weights=[0.5, 0.5], - op=ComparisonOp.GEQ, - bound=0.0, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + scalarized_with_objective_metric = ScalarizedOutcomeConstraint( + metrics=[ + self.metrics["m1"], + self.metrics["m4"], + ], # m1 is objective metric + weights=[0.5, 0.5], + op=ComparisonOp.GEQ, + bound=0.0, + ) with self.assertRaisesRegex( ValueError, "Cannot constrain on objective metric." ): @@ -258,46 +279,59 @@ def setUp(self) -> None: "m2": Metric(name="m2", lower_is_better=False), "m3": Metric(name="m3", lower_is_better=False), } - self.objectives = { - "o1": Objective(metric=self.metrics["m1"]), - "o2": Objective(metric=self.metrics["m2"], minimize=False), - "o3": Objective(metric=self.metrics["m3"], minimize=False), - } - self.objective = Objective(metric=self.metrics["m1"], minimize=True) - self.multi_objective = MultiObjective( - objectives=[self.objectives["o1"], self.objectives["o2"]] - ) - self.scalarized_objective = ScalarizedObjective( - metrics=list(self.metrics.values()), - weights=[-1.0, 1.0, 1.0], - minimize=False, - ) - self.outcome_constraint = OutcomeConstraint( - metric=self.metrics["m3"], op=ComparisonOp.GEQ, bound=-0.25 - ) - self.additional_outcome_constraint = OutcomeConstraint( - metric=self.metrics["m3"], op=ComparisonOp.LEQ, bound=0.25 - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + self.objectives = { + "o1": Objective(metric=self.metrics["m1"]), + "o2": Objective(metric=self.metrics["m2"], minimize=False), + "o3": Objective(metric=self.metrics["m3"], minimize=False), + } + self.objective = Objective(metric=self.metrics["m1"], minimize=True) + self.multi_objective = MultiObjective( + objectives=[self.objectives["o1"], self.objectives["o2"]] + ) + self.scalarized_objective = ScalarizedObjective( + metrics=list(self.metrics.values()), + weights=[-1.0, 1.0, 1.0], + minimize=False, + ) + self.outcome_constraint = OutcomeConstraint( + metric=self.metrics["m3"], op=ComparisonOp.GEQ, bound=-0.25 + ) + self.additional_outcome_constraint = OutcomeConstraint( + metric=self.metrics["m3"], op=ComparisonOp.LEQ, bound=0.25 + ) self.outcome_constraints = [ self.outcome_constraint, self.additional_outcome_constraint, ] - self.objective_thresholds = [ - ObjectiveThreshold(metric=self.metrics["m1"], bound=-1.0, relative=False), - ObjectiveThreshold(metric=self.metrics["m2"], bound=-1.0, relative=False), - ] - self.relative_objective_thresholds = [ - ObjectiveThreshold(metric=self.metrics["m1"], bound=-1.0, relative=True), - ObjectiveThreshold( - metric=self.metrics["m2"], + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + self.objective_thresholds = [ + ObjectiveThreshold( + metric=self.metrics["m1"], bound=-1.0, relative=False + ), + ObjectiveThreshold( + metric=self.metrics["m2"], bound=-1.0, relative=False + ), + ] + self.relative_objective_thresholds = [ + ObjectiveThreshold( + metric=self.metrics["m1"], bound=-1.0, relative=True + ), + ObjectiveThreshold( + metric=self.metrics["m2"], + op=ComparisonOp.GEQ, + bound=-1.0, + relative=True, + ), + ] + self.m3_constraint = OutcomeConstraint( + metric=self.metrics["m3"], op=ComparisonOp.GEQ, - bound=-1.0, + bound=0.1, relative=True, - ), - ] - self.m3_constraint = OutcomeConstraint( - metric=self.metrics["m3"], op=ComparisonOp.GEQ, bound=0.1, relative=True - ) + ) def test_Init(self) -> None: config1 = MultiObjectiveOptimizationConfig( @@ -313,7 +347,7 @@ def test_Init(self) -> None: config1.objective = self.objective # Wrong objective type # updating constraints is fine. config1.outcome_constraints = [self.outcome_constraint] - self.assertEqual(len(config1.metrics), 3) + self.assertEqual(len(config1.metric_names), 3) # objective without outcome_constraints is also supported config2 = MultiObjectiveOptimizationConfig(objective=self.multi_objective) @@ -372,9 +406,11 @@ def test_Eq(self) -> None: ) self.assertEqual(config1, config2) - new_outcome_constraint = OutcomeConstraint( - metric=self.metrics["m3"], op=ComparisonOp.LEQ, bound=0.5 - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + new_outcome_constraint = OutcomeConstraint( + metric=self.metrics["m3"], op=ComparisonOp.LEQ, bound=0.5 + ) config3 = MultiObjectiveOptimizationConfig( objective=self.multi_objective, outcome_constraints=[self.outcome_constraint, new_outcome_constraint], @@ -392,9 +428,14 @@ def test_ConstraintValidation(self) -> None: MultiObjectiveOptimizationConfig(objective=self.objective) # Using an outcome constraint for an objective should raise - outcome_constraint_m1 = OutcomeConstraint( - metric=self.metrics["m1"], op=ComparisonOp.LEQ, bound=1234, relative=False - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + outcome_constraint_m1 = OutcomeConstraint( + metric=self.metrics["m1"], + op=ComparisonOp.LEQ, + bound=1234, + relative=False, + ) with self.assertRaisesRegex( ValueError, "Cannot constrain on objective metric." ): @@ -404,11 +445,13 @@ def test_ConstraintValidation(self) -> None: ) # Two outcome_constraints on the same metric with the same op # should raise. - duplicate_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - op=self.outcome_constraint.op, - bound=self.outcome_constraint.bound + 1, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + duplicate_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + op=self.outcome_constraint.op, + bound=self.outcome_constraint.bound + 1, + ) with self.assertRaises(ValueError): MultiObjectiveOptimizationConfig( objective=self.multi_objective, @@ -416,12 +459,14 @@ def test_ConstraintValidation(self) -> None: ) # Three outcome_constraints on the same metric should raise. - opposing_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. - op=not self.outcome_constraint.op, - bound=self.outcome_constraint.bound, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + opposing_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. + op=not self.outcome_constraint.op, + bound=self.outcome_constraint.bound, + ) with self.assertRaises(ValueError): MultiObjectiveOptimizationConfig( objective=self.multi_objective, @@ -431,12 +476,14 @@ def test_ConstraintValidation(self) -> None: # Two outcome_constraints on the same metric with different ops and # flipped bounds (lower < upper) should raise. add_bound = 1 if self.outcome_constraint.op == ComparisonOp.LEQ else -1 - opposing_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. - op=not self.outcome_constraint.op, - bound=self.outcome_constraint.bound + add_bound, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + opposing_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. + op=not self.outcome_constraint.op, + bound=self.outcome_constraint.bound + add_bound, + ) with self.assertRaises(ValueError): MultiObjectiveOptimizationConfig( objective=self.multi_objective, @@ -445,12 +492,14 @@ def test_ConstraintValidation(self) -> None: # Two outcome_constraints on the same metric with different ops and # bounds should not raise. - opposing_constraint = OutcomeConstraint( - metric=self.outcome_constraint.metric, - # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. - op=not self.outcome_constraint.op, - bound=self.outcome_constraint.bound + 1, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + opposing_constraint = OutcomeConstraint( + metric=self.metrics["m3"], + # pyre-fixme[6]: For 2nd param expected `ComparisonOp` but got `bool`. + op=not self.outcome_constraint.op, + bound=self.outcome_constraint.bound + 1, + ) config = MultiObjectiveOptimizationConfig( objective=self.multi_objective, outcome_constraints=([self.outcome_constraint, opposing_constraint]), @@ -461,12 +510,14 @@ def test_ConstraintValidation(self) -> None: # Test with ScalarizedOutcomeConstraint # should work when not constraining objective - scalarized_constraint = ScalarizedOutcomeConstraint( - metrics=[self.metrics["m3"]], # m3 is not in multi_objective (m1, m2) - weights=[1.0], - op=ComparisonOp.GEQ, - bound=0.0, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + scalarized_constraint = ScalarizedOutcomeConstraint( + metrics=[self.metrics["m3"]], # m3 is not in multi_objective (m1, m2) + weights=[1.0], + op=ComparisonOp.GEQ, + bound=0.0, + ) config_with_scalarized = MultiObjectiveOptimizationConfig( objective=self.multi_objective, outcome_constraints=[scalarized_constraint], @@ -475,15 +526,17 @@ def test_ConstraintValidation(self) -> None: # Can't constrain on metric in ScalarizedOutcomeConstraint # that overlaps with objective - scalarized_with_objective_metric = ScalarizedOutcomeConstraint( - metrics=[ - self.metrics["m1"], - self.metrics["m3"], - ], # m1 is in multi_objective - weights=[0.5, 0.5], - op=ComparisonOp.GEQ, - bound=0.0, - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + scalarized_with_objective_metric = ScalarizedOutcomeConstraint( + metrics=[ + self.metrics["m1"], + self.metrics["m3"], + ], # m1 is in multi_objective + weights=[0.5, 0.5], + op=ComparisonOp.GEQ, + bound=0.0, + ) with self.assertRaisesRegex( ValueError, "Cannot constrain on objective metric." ): @@ -496,13 +549,23 @@ def test_Clone(self) -> None: config1 = MultiObjectiveOptimizationConfig( objective=self.multi_objective, outcome_constraints=self.outcome_constraints ) - self.assertEqual(config1, config1.clone()) + cloned1 = config1.clone() + # Clone normalizes MultiObjective to plain Objective; compare by + # expression and constraints instead of assertEqual. + self.assertEqual(config1.objective.expression, cloned1.objective.expression) + self.assertEqual(config1.outcome_constraints, cloned1.outcome_constraints) + cloned1_moo = assert_is_instance(cloned1, MultiObjectiveOptimizationConfig) + self.assertEqual(config1.objective_thresholds, cloned1_moo.objective_thresholds) config2 = MultiObjectiveOptimizationConfig( objective=self.multi_objective, objective_thresholds=self.objective_thresholds, ) - self.assertEqual(config2, config2.clone()) + cloned2 = config2.clone() + self.assertEqual(config2.objective.expression, cloned2.objective.expression) + self.assertEqual(config2.outcome_constraints, cloned2.outcome_constraints) + cloned2_moo = assert_is_instance(cloned2, MultiObjectiveOptimizationConfig) + self.assertEqual(config2.objective_thresholds, cloned2_moo.objective_thresholds) def test_CloneWithArgs(self) -> None: config1 = MultiObjectiveOptimizationConfig( @@ -514,20 +577,20 @@ def test_CloneWithArgs(self) -> None: objective=self.multi_objective, ) - # Empty args produce exact clone - self.assertEqual( - config1.clone_with_args(), - config1, - ) + # Empty args produce clone with same expression and constraints + cloned = config1.clone_with_args() + self.assertEqual(config1.objective.expression, cloned.objective.expression) + self.assertEqual(config1.outcome_constraints, cloned.outcome_constraints) + self.assertEqual(config1.objective_thresholds, cloned.objective_thresholds) # None args not treated as default - self.assertEqual( - config1.clone_with_args( - outcome_constraints=None, - objective_thresholds=None, - ), - config2, + cloned_none = config1.clone_with_args( + outcome_constraints=None, + objective_thresholds=None, ) + self.assertEqual(config2.objective.expression, cloned_none.objective.expression) + self.assertEqual(cloned_none.outcome_constraints, []) + self.assertEqual(cloned_none.objective_thresholds, []) class PreferenceOptimizationConfigTest(TestCase): @@ -538,14 +601,16 @@ def setUp(self) -> None: "metric2": Metric(name="metric2", lower_is_better=False), "metric3": Metric(name="metric3", lower_is_better=False), } - self.objectives = { - "o1": Objective(metric=self.metrics["metric1"], minimize=True), - "o2": Objective(metric=self.metrics["metric2"], minimize=False), - "o3": Objective(metric=self.metrics["metric3"], minimize=False), - } - self.multi_objective = MultiObjective( - objectives=[self.objectives["o2"], self.objectives["o3"]] - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + self.objectives = { + "o1": Objective(metric=self.metrics["metric1"], minimize=True), + "o2": Objective(metric=self.metrics["metric2"], minimize=False), + "o3": Objective(metric=self.metrics["metric3"], minimize=False), + } + self.multi_objective = MultiObjective( + objectives=[self.objectives["o2"], self.objectives["o3"]] + ) self.preference_profile_name = "pe_exp" def test_Init(self) -> None: @@ -562,17 +627,19 @@ def test_Init(self) -> None: with self.assertRaisesRegex( NotImplementedError, "Outcome constraints are not yet supported" ): - PreferenceOptimizationConfig( - objective=self.multi_objective, - preference_profile_name=self.preference_profile_name, - outcome_constraints=[ - OutcomeConstraint( - metric=self.metrics["metric1"], - op=ComparisonOp.LEQ, - bound=0.5, - ) - ], - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + PreferenceOptimizationConfig( + objective=self.multi_objective, + preference_profile_name=self.preference_profile_name, + outcome_constraints=[ + OutcomeConstraint( + metric=self.metrics["metric1"], + op=ComparisonOp.LEQ, + bound=0.5, + ) + ], + ) def test_Eq(self) -> None: config1 = PreferenceOptimizationConfig( @@ -593,9 +660,11 @@ def test_Eq(self) -> None: self.assertNotEqual(config1, config3) # Different objective - different_objective = MultiObjective( - objectives=[self.objectives["o1"], self.objectives["o2"]] - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + different_objective = MultiObjective( + objectives=[self.objectives["o1"], self.objectives["o2"]] + ) config4 = PreferenceOptimizationConfig( objective=different_objective, preference_profile_name=self.preference_profile_name, @@ -608,12 +677,15 @@ def test_Clone(self) -> None: preference_profile_name=self.preference_profile_name, ) cloned_config = assert_is_instance(config.clone(), PreferenceOptimizationConfig) - self.assertEqual(config, cloned_config) self.assertIsNot(config, cloned_config) self.assertEqual( cloned_config.preference_profile_name, self.preference_profile_name ) - self.assertEqual(cloned_config.objective, self.multi_objective) + # Clone normalizes MultiObjective to plain Objective; compare + # by expression instead of assertEqual. + self.assertEqual( + config.objective.expression, cloned_config.objective.expression + ) config = PreferenceOptimizationConfig( objective=self.multi_objective, @@ -621,19 +693,29 @@ def test_Clone(self) -> None: ) # ======= Clone with args ======= - # Empty args produce exact clone + # Empty args produce clone with matching properties cloned_config = config.clone_with_args() - self.assertEqual(config, cloned_config) self.assertIsNot(config, cloned_config) + self.assertEqual( + config.objective.expression, cloned_config.objective.expression + ) + self.assertEqual( + cloned_config.preference_profile_name, self.preference_profile_name + ) # Clone with different objective - different_objective = MultiObjective( - objectives=[self.objectives["o1"], self.objectives["o3"]] - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + different_objective = MultiObjective( + objectives=[self.objectives["o1"], self.objectives["o3"]] + ) cloned_with_diff_objective = config.clone_with_args( objective=different_objective ) - self.assertEqual(cloned_with_diff_objective.objective, different_objective) + self.assertEqual( + cloned_with_diff_objective.objective.expression, + different_objective.expression, + ) self.assertEqual( cloned_with_diff_objective.preference_profile_name, self.preference_profile_name, @@ -644,7 +726,10 @@ def test_Clone(self) -> None: cloned_with_diff_profile = config.clone_with_args( preference_profile_name=different_profile ) - self.assertEqual(cloned_with_diff_profile.objective, self.multi_objective) + self.assertEqual( + cloned_with_diff_profile.objective.expression, + self.multi_objective.expression, + ) self.assertEqual( cloned_with_diff_profile.preference_profile_name, different_profile ) diff --git a/ax/core/tests/test_outcome_constraint.py b/ax/core/tests/test_outcome_constraint.py index f288ece1703..3d75ce9b8ba 100644 --- a/ax/core/tests/test_outcome_constraint.py +++ b/ax/core/tests/test_outcome_constraint.py @@ -6,102 +6,157 @@ # pyre-strict -from unittest import mock +import warnings from ax.core.metric import Metric from ax.core.outcome_constraint import ( - CONSTRAINT_WARNING_MESSAGE, - LOWER_BOUND_MISMATCH, ObjectiveThreshold, OutcomeConstraint, ScalarizedOutcomeConstraint, - UPPER_BOUND_MISMATCH, ) from ax.core.types import ComparisonOp +from ax.exceptions.core import UserInputError from ax.utils.common.testutils import TestCase -OUTCOME_CONSTRAINT_PATH = "ax.core.outcome_constraint" - - class OutcomeConstraintTest(TestCase): def setUp(self) -> None: super().setUp() - self.minimize_metric = Metric(name="bar", lower_is_better=True) - self.maximize_metric = Metric(name="baz", lower_is_better=False) - self.bound = 0 - simple_metric = Metric(name="foo") - self.constraint = OutcomeConstraint( - metric=simple_metric, op=ComparisonOp.GEQ, bound=self.bound - ) + self.constraint = OutcomeConstraint(expression="foo >= 0", relative=False) + + def test_ExpressionInit(self) -> None: + """Test creating constraints with expression strings.""" + # GEQ constraint + oc = OutcomeConstraint(expression="qps >= 700", relative=False) + self.assertEqual(oc.expression, "qps >= 700") + self.assertEqual(oc.metric_names, ["qps"]) + self.assertEqual(oc.op, ComparisonOp.GEQ) + self.assertEqual(oc.bound, 700.0) + self.assertFalse(oc.relative) + + # LEQ constraint + oc2 = OutcomeConstraint(expression="loss <= 0.5", relative=False) + self.assertEqual(oc2.metric_names, ["loss"]) + self.assertEqual(oc2.op, ComparisonOp.LEQ) + self.assertEqual(oc2.bound, 0.5) + self.assertFalse(oc2.relative) + + # Relative constraint via * baseline + oc3 = OutcomeConstraint(expression="latency <= 1.05 * baseline") + self.assertEqual(oc3.metric_names, ["latency"]) + self.assertEqual(oc3.bound, 5.0) # (1.05 - 1) * 100 + self.assertTrue(oc3.relative) + + # relative kwarg without baseline in expression has no effect + oc4 = OutcomeConstraint(expression="qps >= 700", relative=True) + self.assertFalse(oc4.relative) # expression has no baseline + + # No expression should error + with self.assertRaisesRegex(UserInputError, "expression string is required"): + OutcomeConstraint() + with self.assertRaisesRegex(UserInputError, "expression string is required"): + OutcomeConstraint(expression=None) + + # Missing operator should error + with self.assertRaisesRegex(UserInputError, "Expected an inequality"): + OutcomeConstraint(expression="qps 700") + + def test_DeprecatedMetricInit(self) -> None: + """Test the deprecated metric-based init path.""" + metric = Metric(name="foo") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + oc = OutcomeConstraint( + metric=metric, op=ComparisonOp.GEQ, bound=0, relative=True + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + self.assertTrue(len(deprecation_warnings) > 0) + + self.assertEqual(oc.metric_names, ["foo"]) + self.assertEqual(oc.op, ComparisonOp.GEQ) + self.assertEqual(oc.bound, 0.0) + self.assertTrue(oc.relative) + + # Cannot specify both expression and metric + with self.assertRaisesRegex(UserInputError, "Cannot specify both"): + OutcomeConstraint( + expression="foo >= 0", metric=metric, op=ComparisonOp.GEQ, bound=0 + ) + + # Must specify op and bound with metric + with self.assertRaisesRegex(UserInputError, "Must specify `op`"): + OutcomeConstraint(metric=metric, bound=0) + with self.assertRaisesRegex(UserInputError, "Must specify `bound`"): + OutcomeConstraint(metric=metric, op=ComparisonOp.GEQ) def test_Eq(self) -> None: - constraint1 = OutcomeConstraint( - metric=Metric(name="foo"), op=ComparisonOp.GEQ, bound=self.bound - ) - constraint2 = OutcomeConstraint( - metric=Metric(name="foo"), op=ComparisonOp.GEQ, bound=self.bound - ) + constraint1 = OutcomeConstraint(expression="foo >= 0", relative=False) + constraint2 = OutcomeConstraint(expression="foo >= 0", relative=False) self.assertEqual(constraint1, constraint2) - constraint3 = OutcomeConstraint( - metric=Metric(name="foo"), op=ComparisonOp.LEQ, bound=self.bound - ) + constraint3 = OutcomeConstraint(expression="foo <= 0", relative=False) self.assertNotEqual(constraint1, constraint3) - def test_ValidMutations(self) -> None: - # updating constraint metric is ok as long as lower_is_better is compatible. - self.constraint.metric = self.maximize_metric - self.constraint.op = ComparisonOp.LEQ - self.assertEqual(self.constraint.metric.name, "baz") - self.assertEqual(self.constraint.metric.signature, "baz") - - def test_OutcomeConstraintFail(self) -> None: - logger_name = OUTCOME_CONSTRAINT_PATH + ".logger" - with mock.patch(logger_name) as mock_warning: - OutcomeConstraint( - metric=self.minimize_metric, op=ComparisonOp.GEQ, bound=self.bound - ) - mock_warning.debug.assert_called_once_with( - CONSTRAINT_WARNING_MESSAGE.format(**LOWER_BOUND_MISMATCH, name="bar") - ) - with mock.patch(logger_name) as mock_warning: - OutcomeConstraint( - metric=self.maximize_metric, op=ComparisonOp.LEQ, bound=self.bound - ) - mock_warning.debug.assert_called_once_with( - CONSTRAINT_WARNING_MESSAGE.format(**UPPER_BOUND_MISMATCH, name="baz") - ) - def test_Sortable(self) -> None: - constraint1 = OutcomeConstraint( - metric=Metric(name="foo"), op=ComparisonOp.LEQ, bound=self.bound - ) - constraint2 = OutcomeConstraint( - metric=Metric(name="foo"), op=ComparisonOp.GEQ, bound=self.bound - ) + constraint1 = OutcomeConstraint(expression="foo <= 0", relative=False) + constraint2 = OutcomeConstraint(expression="foo >= 0", relative=False) self.assertTrue(constraint1 < constraint2) - def test_validate_constraint(self) -> None: - metric = Metric(name="metric0", lower_is_better=False) - oc = OutcomeConstraint(metric, bound=-3, relative=True, op=ComparisonOp.GEQ) - self.assertTrue(oc._validate_constraint()[0]) - oc = OutcomeConstraint(metric, bound=-3, relative=True, op=ComparisonOp.LEQ) - self.assertFalse(oc._validate_constraint()[0]) - oc = OutcomeConstraint(metric, bound=3, relative=True, op=ComparisonOp.GEQ) - self.assertFalse(oc._validate_constraint()[0]) - oc = OutcomeConstraint(metric, bound=3, relative=True, op=ComparisonOp.LEQ) - self.assertFalse(oc._validate_constraint()[0]) - - metric = Metric(name="metric1", lower_is_better=True) - oc = OutcomeConstraint(metric, bound=3, relative=True, op=ComparisonOp.LEQ) - self.assertTrue(oc._validate_constraint()[0]) - oc = OutcomeConstraint(metric, bound=3, relative=True, op=ComparisonOp.GEQ) - self.assertFalse(oc._validate_constraint()[0]) - oc = OutcomeConstraint(metric, bound=-3, relative=True, op=ComparisonOp.LEQ) - self.assertFalse(oc._validate_constraint()[0]) - oc = OutcomeConstraint(metric, bound=-3, relative=True, op=ComparisonOp.GEQ) - self.assertFalse(oc._validate_constraint()[0]) + def test_Clone(self) -> None: + oc = OutcomeConstraint(expression="foo >= 10", relative=False) + cloned = oc.clone() + self.assertEqual(oc, cloned) + self.assertIsNot(oc, cloned) + self.assertEqual(oc.expression, cloned.expression) + + def test_Repr(self) -> None: + oc = OutcomeConstraint(expression="foo >= 10", relative=False) + self.assertEqual(str(oc), "OutcomeConstraint(foo >= 10)") + + oc_rel = OutcomeConstraint(expression="bar <= 5 * baseline") + self.assertEqual(str(oc_rel), "OutcomeConstraint(bar <= 5 * baseline)") + + def test_DeprecatedMetricProperty(self) -> None: + oc = OutcomeConstraint(expression="foo >= 10", relative=False) + + # Use a helper to prevent Pyre from resolving the attribute name + # statically; the .metric attribute was intentionally removed. + def _getattr(obj: object, name: str) -> object: + return getattr(obj, name) + + with self.assertRaises(AttributeError): + _getattr(oc, "metric") + + def test_ScalarizedExpression(self) -> None: + """Test a constraint with multiple weighted metrics.""" + oc = OutcomeConstraint(expression="2*m1 + 3*m2 <= 10", relative=False) + self.assertEqual(sorted(oc.metric_names), ["m1", "m2"]) + self.assertEqual(oc.op, ComparisonOp.LEQ) + self.assertEqual(oc.bound, 10.0) + weights_dict = dict(oc.metric_weights) + self.assertEqual(weights_dict["m1"], 2.0) + self.assertEqual(weights_dict["m2"], 3.0) + + def test_RelativeRoundTrip(self) -> None: + """Test that relative constraints round-trip correctly.""" + # GEQ relative constraint + oc = OutcomeConstraint(expression="qps >= 0.95 * baseline") + self.assertEqual(oc.bound, -5.0) + self.assertEqual(oc.op, ComparisonOp.GEQ) + self.assertTrue(oc.relative) + + # LEQ relative constraint + oc2 = OutcomeConstraint(expression="latency <= 1.05 * baseline") + self.assertEqual(oc2.bound, 5.0) + self.assertEqual(oc2.op, ComparisonOp.LEQ) + self.assertTrue(oc2.relative) + + # Standalone baseline (multiplier = 1) + oc3 = OutcomeConstraint(expression="qps >= baseline") + self.assertEqual(oc3.bound, 0.0) + self.assertTrue(oc3.relative) class ObjectiveThresholdTest(TestCase): @@ -110,150 +165,229 @@ def setUp(self) -> None: self.minimize_metric = Metric(name="bar", lower_is_better=True) self.maximize_metric = Metric(name="baz", lower_is_better=False) self.ambiguous_metric = Metric(name="buz") - self.bound = 0 - self.threshold = ObjectiveThreshold( - metric=self.maximize_metric, op=ComparisonOp.GEQ, bound=self.bound - ) + + def test_Init(self) -> None: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + ot = ObjectiveThreshold( + metric=self.maximize_metric, + bound=0.5, + relative=False, + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + self.assertTrue(len(deprecation_warnings) > 0) + + # GEQ inferred from lower_is_better=False + self.assertEqual(ot.op, ComparisonOp.GEQ) + self.assertEqual(ot.bound, 0.5) + self.assertFalse(ot.relative) + self.assertEqual(ot.metric_names, ["baz"]) + + def test_OpInference(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot_min = ObjectiveThreshold(metric=self.minimize_metric, bound=0.0) + self.assertEqual(ot_min.op, ComparisonOp.LEQ) + + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot_max = ObjectiveThreshold(metric=self.maximize_metric, bound=0.0) + self.assertEqual(ot_max.op, ComparisonOp.GEQ) + + def test_OpManualOverride(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot = ObjectiveThreshold( + metric=self.maximize_metric, + bound=0.0, + op=ComparisonOp.LEQ, + ) + self.assertEqual(ot.op, ComparisonOp.LEQ) + + def test_AmbiguousMetricRequiresOp(self) -> None: + with self.assertRaises(ValueError): + ObjectiveThreshold(metric=self.ambiguous_metric, bound=0.0) def test_Eq(self) -> None: - threshold1 = ObjectiveThreshold(metric=self.minimize_metric, bound=self.bound) - threshold2 = ObjectiveThreshold(metric=self.minimize_metric, bound=self.bound) + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + threshold1 = ObjectiveThreshold(metric=self.minimize_metric, bound=0.0) + threshold2 = ObjectiveThreshold(metric=self.minimize_metric, bound=0.0) self.assertEqual(threshold1, threshold2) - threshold3 = ObjectiveThreshold( - metric=self.minimize_metric, bound=self.bound, relative=False - ) + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + threshold3 = ObjectiveThreshold( + metric=self.minimize_metric, bound=0.0, relative=False + ) self.assertNotEqual(threshold1, threshold3) - constraint3 = OutcomeConstraint( - metric=self.minimize_metric, op=ComparisonOp.LEQ, bound=self.bound - ) - self.assertNotEqual(threshold1, constraint3) - - def test_ValidMutations(self) -> None: - # updating constraint metric is ok as long as lower_is_better is compatible. - self.threshold.metric = self.ambiguous_metric - self.threshold.op = ComparisonOp.LEQ - self.assertEqual(self.threshold.metric.name, "buz") - self.assertEqual(self.threshold.metric.signature, "buz") - - def test_ObjectiveThresholdFail(self) -> None: - logger_name = OUTCOME_CONSTRAINT_PATH + ".logger" - with mock.patch(logger_name) as mock_warning: - ObjectiveThreshold( - metric=self.minimize_metric, op=ComparisonOp.GEQ, bound=self.bound + def test_Clone(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot = ObjectiveThreshold( + metric=self.maximize_metric, bound=0.5, relative=False ) - mock_warning.debug.assert_called_once_with( - CONSTRAINT_WARNING_MESSAGE.format(**LOWER_BOUND_MISMATCH, name="bar") - ) - with mock.patch(logger_name) as mock_warning: - ObjectiveThreshold( - metric=self.maximize_metric, op=ComparisonOp.LEQ, bound=self.bound + cloned = ot.clone() + self.assertEqual(ot, cloned) + self.assertIsNot(ot, cloned) + self.assertIsInstance(cloned, ObjectiveThreshold) + + def test_Repr(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot = ObjectiveThreshold( + metric=self.maximize_metric, + bound=0.5, + relative=False, + op=ComparisonOp.GEQ, ) - mock_warning.debug.assert_called_once_with( - CONSTRAINT_WARNING_MESSAGE.format(**UPPER_BOUND_MISMATCH, name="baz") - ) + self.assertEqual(str(ot), "ObjectiveThreshold(baz >= 0.5)") def test_Relativize(self) -> None: - self.assertTrue( - ObjectiveThreshold( - metric=self.maximize_metric, op=ComparisonOp.LEQ, bound=self.bound - ).relative - ) - self.assertTrue( - ObjectiveThreshold( + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + # Default is relative=True + ot_default = ObjectiveThreshold( metric=self.maximize_metric, op=ComparisonOp.LEQ, - bound=self.bound, - relative=True, - ).relative - ) - self.assertFalse( - ObjectiveThreshold( + bound=0.0, + ) + self.assertTrue(ot_default.relative) + + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot_abs = ObjectiveThreshold( metric=self.maximize_metric, op=ComparisonOp.LEQ, - bound=self.bound, + bound=0.0, relative=False, - ).relative - ) + ) + self.assertFalse(ot_abs.relative) + + def test_IsInstance(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + ot = ObjectiveThreshold( + metric=self.maximize_metric, bound=0.5, relative=False + ) + self.assertIsInstance(ot, OutcomeConstraint) + self.assertIsInstance(ot, ObjectiveThreshold) class ScalarizedOutcomeConstraintTest(TestCase): def setUp(self) -> None: super().setUp() self.metrics = [ - Metric(name="m1", lower_is_better=True), - Metric(name="m2", lower_is_better=True), - Metric(name="m3", lower_is_better=True), + Metric(name="m1"), + Metric(name="m2"), + Metric(name="m3"), ] self.weights = [0.1, 0.3, 0.6] - self.bound = 0 - self.constraint = ScalarizedOutcomeConstraint( - metrics=self.metrics, - weights=self.weights, - op=ComparisonOp.GEQ, - bound=self.bound, - ) def test_Init(self) -> None: - self.assertListEqual(self.constraint.metrics, self.metrics) - self.assertListEqual(self.constraint.weights, self.weights) - self.assertEqual(len(list(self.constraint.metric_weights)), len(self.metrics)) - self.assertEqual( - str(self.constraint), - ("ScalarizedOutcomeConstraint(0.1 * m1 + 0.3 * m2 + 0.6 * m3 >= 0%)"), - ) - # check that weights are set uniformly by default - con = ScalarizedOutcomeConstraint( - metrics=[ - Metric(name="m1", lower_is_better=True), - Metric(name="m2", lower_is_better=True), - ], - op=ComparisonOp.LEQ, - bound=self.bound, - ) - self.assertListEqual(con.weights, [0.5, 0.5]) - - def test_Eq(self) -> None: - constraint1 = ScalarizedOutcomeConstraint( - metrics=self.metrics, - weights=self.weights, - op=ComparisonOp.GEQ, - bound=self.bound, - ) - self.assertEqual(constraint1, self.constraint) - - constraint2 = ScalarizedOutcomeConstraint( - metrics=self.metrics, - weights=[0.2, 0.2, 0.6], - op=ComparisonOp.LEQ, - bound=self.bound, - ) - self.assertNotEqual(constraint2, self.constraint) - - def test_Clone(self) -> None: - self.assertEqual(self.constraint, self.constraint.clone()) - - def test_ValidMutations(self) -> None: - # updating constraint metric is ok as long as lower_is_better is compatible. - self.constraint.metrics = [ - Metric(name="m2"), - Metric(name="m4"), - ] - self.constraint.op = ComparisonOp.LEQ + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + soc = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=self.weights, + op=ComparisonOp.GEQ, + bound=0, + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + self.assertTrue(len(deprecation_warnings) > 0) + + self.assertEqual(soc.weights, self.weights) + self.assertEqual(sorted(soc.metric_names), ["m1", "m2", "m3"]) + self.assertEqual(soc.op, ComparisonOp.GEQ) + self.assertEqual(soc.bound, 0.0) + + def test_DefaultWeights(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + soc = ScalarizedOutcomeConstraint( + metrics=[Metric(name="m1"), Metric(name="m2")], + op=ComparisonOp.LEQ, + bound=0, + ) + self.assertEqual(soc.weights, [0.5, 0.5]) - def test_RaiseError(self) -> None: - # set a wrong weights + def test_WeightMismatchRaises(self) -> None: with self.assertRaises(ValueError): ScalarizedOutcomeConstraint( metrics=self.metrics, weights=[0.2, 0.8], op=ComparisonOp.LEQ, - bound=self.bound, + bound=0, ) - with self.assertRaises(NotImplementedError): - self.constraint.metric + def test_Eq(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + soc1 = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=self.weights, + op=ComparisonOp.GEQ, + bound=0, + ) + soc2 = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=self.weights, + op=ComparisonOp.GEQ, + bound=0, + ) + self.assertEqual(soc1, soc2) + + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + soc3 = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=[0.2, 0.2, 0.6], + op=ComparisonOp.LEQ, + bound=0, + ) + self.assertNotEqual(soc1, soc3) - with self.assertRaises(NotImplementedError): - self.constraint.metric = self.metrics[0] + def test_Clone(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + soc = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=self.weights, + op=ComparisonOp.GEQ, + bound=0, + ) + cloned = soc.clone() + self.assertEqual(soc, cloned) + self.assertIsNot(soc, cloned) + self.assertIsInstance(cloned, ScalarizedOutcomeConstraint) + self.assertEqual(cloned.weights, self.weights) + + def test_Repr(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + soc = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=self.weights, + op=ComparisonOp.GEQ, + bound=0, + ) + repr_str = str(soc) + self.assertTrue(repr_str.startswith("ScalarizedOutcomeConstraint(")) + self.assertIn(">=", repr_str) + + def test_IsInstance(self) -> None: + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + soc = ScalarizedOutcomeConstraint( + metrics=self.metrics, + weights=self.weights, + op=ComparisonOp.GEQ, + bound=0, + ) + self.assertIsInstance(soc, OutcomeConstraint) + self.assertIsInstance(soc, ScalarizedOutcomeConstraint) diff --git a/ax/core/tests/test_trial.py b/ax/core/tests/test_trial.py index d638968896d..0bfd29d27f2 100644 --- a/ax/core/tests/test_trial.py +++ b/ax/core/tests/test_trial.py @@ -41,11 +41,11 @@ [ { "arm_name": "0_0", - "metric_name": get_objective().metric.name, + "metric_name": get_objective().metric_names[0], "mean": 1.0, "sem": 2.0, "trial_index": 0, - "metric_signature": get_objective().metric.signature, + "metric_signature": get_objective().metric_names[0], } ] ) @@ -324,10 +324,10 @@ def test_objective_mean_empty_df(self, _mock: Mock) -> None: @patch( f"{BaseTrial.__module__}.{BaseTrial.__name__}.fetch_data_results", - return_value={get_objective().metric.signature: Ok(TEST_DATA)}, + return_value={get_objective().metric_names[0]: Ok(TEST_DATA)}, ) def test_fetch_data_result(self, mock: Mock) -> None: - metric_name = get_objective().metric.name + metric_name = get_objective().metric_names[0] results = self.experiment.trials[0].fetch_data_results() self.assertTrue(results[metric_name].is_ok()) diff --git a/ax/core/tests/test_utils.py b/ax/core/tests/test_utils.py index 83208838ac2..4659e5095ea 100644 --- a/ax/core/tests/test_utils.py +++ b/ax/core/tests/test_utils.py @@ -264,7 +264,7 @@ def test_update_trial_status(self) -> None: } for metric_name in none_throws( experiment.optimization_config - ).metrics + ).metric_names ] ) ) @@ -346,7 +346,7 @@ def test_completed_complete_trials_not_pending(self) -> None: } for metric_name in none_throws( experiment.optimization_config - ).metrics + ).metric_names ] ) ) @@ -361,14 +361,14 @@ def test_completed_complete_trials_not_pending(self) -> None: # NOT be pending for any metric. pending = get_pending_observation_features(experiment) if pending is not None: - for metric_name in none_throws(experiment.optimization_config).metrics: + for metric_name in none_throws(experiment.optimization_config).metric_names: self.assertNotIn(obs_feat, pending.get(metric_name, [])) # get_pending_observation_features_based_on_trial_status: # COMPLETED+COMPLETE trial should NOT be pending. pending_status = get_pending_status(experiment) if pending_status is not None: - for metric_name in none_throws(experiment.optimization_config).metrics: + for metric_name in none_throws(experiment.optimization_config).metric_names: self.assertNotIn(obs_feat, pending_status.get(metric_name, [])) def test_get_pending_observation_features_multi_trial(self) -> None: @@ -1025,6 +1025,8 @@ def test_curve_data(self) -> None: with_trial=True, with_completed_trial=False, ) + exp.add_metric(Metric(name="metric_a")) + exp.add_metric(Metric(name="metric_b")) exp.optimization_config = OptimizationConfig( objective=Objective(metric=Metric(name="metric_a"), minimize=False), outcome_constraints=[ @@ -1067,6 +1069,8 @@ def test_curve_data(self) -> None: with_trial=True, with_completed_trial=False, ) + exp2.add_metric(Metric(name="metric_a")) + exp2.add_metric(Metric(name="metric_b")) exp2.optimization_config = none_throws(exp.optimization_config) trial2 = exp2.trials[0] trial2.mark_running(no_runner_required=True) diff --git a/ax/core/trial.py b/ax/core/trial.py index 235ad683a46..0f70e170f19 100644 --- a/ax/core/trial.py +++ b/ax/core/trial.py @@ -197,7 +197,7 @@ def objective_mean(self) -> float: raise ValueError( "Experiment optimization config (and thus the objective) is not set." ) - return self.get_metric_mean(metric_name=opt_config.objective.metric.name) + return self.get_metric_mean(metric_name=opt_config.objective.metric_names[0]) def get_metric_mean(self, metric_name: str) -> float: """Metric mean for the arm attached to this trial, retrieved from the diff --git a/ax/core/utils.py b/ax/core/utils.py index e0e8b298f46..994fa1dfe16 100644 --- a/ax/core/utils.py +++ b/ax/core/utils.py @@ -6,7 +6,7 @@ # pyre-strict -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence from copy import deepcopy from datetime import datetime, timedelta from enum import Enum @@ -21,6 +21,7 @@ from ax.core.experiment import Experiment from ax.core.generator_run import GeneratorRun from ax.core.map_metric import MapMetric +from ax.core.metric import Metric from ax.core.observation import ObservationFeatures from ax.core.optimization_config import OptimizationConfig from ax.core.trial import Trial @@ -57,7 +58,7 @@ def _maybe_update_trial_status_to_complete( if experiment.optimization_config is not None: optimization_config = experiment.optimization_config trial_data = experiment.lookup_data(trial_indices=[trial_index]) - missing_metrics = set(optimization_config.metrics.keys()) - { + missing_metrics = {*optimization_config.metric_names} - { *trial_data.metric_names } @@ -206,7 +207,7 @@ def compute_metric_availability( "Either pass one explicitly, set one on the experiment, or provide " "metric_names." ) - required_metrics = set(resolved_opt_config.metrics.keys()) + required_metrics = resolved_opt_config.metric_names # Resolve trial indices. if trial_indices is not None: @@ -731,18 +732,15 @@ def _time_trial_completed_safe(trial: BatchTrial) -> datetime: # -------------------- MapMetric related utils. --------------------- -def has_map_metrics(optimization_config: OptimizationConfig) -> bool: - """Check if the optimization config has any ``MapMetric``s. +def has_map_metrics(metrics: Sequence[Metric]) -> bool: + """Check if the given metrics contain any ``MapMetric``s. Args: - optimization_config: Optimization config. + metrics: A mapping from metric name to Metric object. """ - metrics = optimization_config.metrics - # Technically an OptimizationConfig could have zero metrics since a - # MultiObjective could have 0 objectives if len(metrics) == 0: return False - return any(isinstance(metric, MapMetric) for metric in metrics.values()) + return any(isinstance(metric, MapMetric) for metric in metrics) # -------------------- Context manager and decorator utils. --------------------- diff --git a/ax/early_stopping/dispatch.py b/ax/early_stopping/dispatch.py index 8c3a68423b7..c71f285ecb1 100644 --- a/ax/early_stopping/dispatch.py +++ b/ax/early_stopping/dispatch.py @@ -35,7 +35,10 @@ def get_default_ess_or_none( opt_config is None or isinstance(opt_config, MultiObjectiveOptimizationConfig) or len(opt_config.outcome_constraints) > 0 - or not any(isinstance(m, MapMetric) for m in opt_config.metrics.values()) + or not any( + isinstance(experiment.get_metric(name), MapMetric) + for name in opt_config.metric_names + ) ): return None diff --git a/ax/early_stopping/experiment_replay.py b/ax/early_stopping/experiment_replay.py index d4a1f5a5fe2..a6b81b7226e 100644 --- a/ax/early_stopping/experiment_replay.py +++ b/ax/early_stopping/experiment_replay.py @@ -89,6 +89,7 @@ def replay_experiment( optimization_config=optimization_config, search_space=dummy_search_space, runner=runner, + metrics=[replay_metric], ) # Setup a Orchestrator with a dummy gs to replay the historical experiment diff --git a/ax/early_stopping/strategies/base.py b/ax/early_stopping/strategies/base.py index 15c340fb72c..c3671a55dbe 100644 --- a/ax/early_stopping/strategies/base.py +++ b/ax/early_stopping/strategies/base.py @@ -17,7 +17,6 @@ from ax.core.batch_trial import BatchTrial from ax.core.data import Data, MAP_KEY from ax.core.experiment import Experiment -from ax.core.objective import MultiObjective from ax.core.trial_status import TrialStatus from ax.early_stopping.utils import ( _interval_boundary, @@ -505,8 +504,8 @@ def _all_objectives_and_directions(self, experiment: Experiment) -> dict[str, bo Args: experiment: The experiment containing the optimization config. - Returns: A dictionary mapping metric names to a Boolean indicating whether - the objective is being minimized. + Returns: A dictionary mapping metric signatures to a Boolean indicating + whether the objective is being minimized. """ if self.metric_signatures is None: logger.debug( @@ -516,15 +515,15 @@ def _all_objectives_and_directions(self, experiment: Experiment) -> dict[str, bo ) optimization_config = none_throws(experiment.optimization_config) objective = optimization_config.objective - objectives = ( - objective.objectives - if isinstance(objective, MultiObjective) - else [objective] - ) directions = {} - for objective in objectives: - metric_signature = objective.metric.signature - directions[metric_signature] = objective.minimize + if objective.is_multi_objective: + for name, weight in objective.metric_weights: + metric = experiment.get_metric(name) + directions[metric.signature] = weight < 0 + else: + name = objective.metric_names[0] + metric = experiment.get_metric(name) + directions[metric.signature] = objective.minimize else: metric_signatures = list(self.metric_signatures) diff --git a/ax/early_stopping/tests/test_experiment_replay.py b/ax/early_stopping/tests/test_experiment_replay.py index a0a9dc0d67c..f62873dd27b 100644 --- a/ax/early_stopping/tests/test_experiment_replay.py +++ b/ax/early_stopping/tests/test_experiment_replay.py @@ -25,13 +25,17 @@ def setUp(self) -> None: super().setUp() # Experiment with MapMetric for tests that need a valid default ESS. self.exp = get_branin_experiment_with_timestamp_map_metric() - self.metric = none_throws(self.exp.optimization_config).objective.metric + metric_name = none_throws(self.exp.optimization_config).objective.metric_names[ + 0 + ] + self.metric = self.exp.get_metric(metric_name) def test_estimate_hypothetical_ess_no_default_strategy(self) -> None: """Test that UnsupportedError is raised when no default ESS is available.""" # Non-MapMetric experiment has no default ESS. exp = get_branin_experiment(has_optimization_config=True) - metric = none_throws(exp.optimization_config).objective.metric + metric_name = none_throws(exp.optimization_config).objective.metric_names[0] + metric = exp.get_metric(metric_name) with self.assertRaises(UnsupportedError) as e: estimate_hypothetical_early_stopping_savings( diff --git a/ax/early_stopping/tests/test_strategies.py b/ax/early_stopping/tests/test_strategies.py index e5751fdb946..45380fac071 100644 --- a/ax/early_stopping/tests/test_strategies.py +++ b/ax/early_stopping/tests/test_strategies.py @@ -296,7 +296,7 @@ def test_default_objective_and_direction(self, _: MagicMock) -> None: test_objective = none_throws(test_experiment.optimization_config).objective with self.subTest("provide metric names"): es_strategy = FakeStrategy( - metric_signatures=[test_objective.metric.signature] + metric_signatures=[test_objective.metric_names[0]] ) ( actual_metric_name, @@ -305,7 +305,7 @@ def test_default_objective_and_direction(self, _: MagicMock) -> None: self.assertEqual( actual_metric_name, - test_objective.metric.name, + test_objective.metric_names[0], ) self.assertEqual( actual_minimize, @@ -322,7 +322,7 @@ def test_default_objective_and_direction(self, _: MagicMock) -> None: self.assertEqual( actual_metric_name, - test_objective.metric.name, + test_objective.metric_names[0], ) self.assertEqual( actual_minimize, @@ -344,16 +344,16 @@ def test_default_objective_and_direction(self, _: MagicMock) -> None: ) self.assertEqual( actual_metric_name, - test_multi_objective.objectives[0].metric.name, + test_multi_objective.metric_names[0], ) self.assertEqual( actual_minimize, - test_multi_objective.objectives[0].minimize, + test_multi_objective.metric_weights[0][1] < 0, ) with self.subTest("provide metric names -- multi-objective"): es_strategy = FakeStrategy( - metric_signatures=[test_multi_objective.objectives[1].metric.signature] + metric_signatures=[test_multi_objective.metric_names[1]] ) ( actual_metric_name, @@ -363,11 +363,11 @@ def test_default_objective_and_direction(self, _: MagicMock) -> None: ) self.assertEqual( actual_metric_name, - test_multi_objective.objectives[1].metric.name, + test_multi_objective.metric_names[1], ) self.assertEqual( actual_minimize, - test_multi_objective.objectives[1].minimize, + test_multi_objective.metric_weights[1][1] < 0, ) @patch.object(logger, "warning") @@ -852,9 +852,7 @@ def _test_percentile_early_stopping_strategy( # be used for early stopping exp._optimization_config = None data = exp.fetch_data() - self.assertTrue( - (data.full_df["metric_name"] == "tracking_branin_map").all() - ) + self.assertIn("tracking_branin_map", data.full_df["metric_name"].values) else: metric_signatures = None diff --git a/ax/generation_strategy/tests/test_dispatch_utils.py b/ax/generation_strategy/tests/test_dispatch_utils.py index 9351d41d3d0..707b06a847c 100644 --- a/ax/generation_strategy/tests/test_dispatch_utils.py +++ b/ax/generation_strategy/tests/test_dispatch_utils.py @@ -13,7 +13,7 @@ import torch from ax.adapter.registry import Generators from ax.adapter.transforms.log_y import LogY -from ax.core.objective import MultiObjective +from ax.core.objective import Objective from ax.core.optimization_config import MultiObjectiveOptimizationConfig from ax.generation_strategy.dispatch_utils import ( _make_botorch_step, @@ -166,7 +166,7 @@ def test_choose_generation_strategy_legacy(self) -> None: ) with self.subTest("MOO"): optimization_config = MultiObjectiveOptimizationConfig( - objective=MultiObjective(objectives=[]) + objective=Objective(expression="branin_a, branin_b") ) sobol_gpei = choose_generation_strategy_legacy( search_space=get_branin_search_space(), @@ -337,7 +337,7 @@ def test_choose_generation_strategy_legacy(self) -> None: search_space = get_branin_search_space(with_choice_parameter=True) search_space.parameters["x2"]._is_ordered = False optimization_config = MultiObjectiveOptimizationConfig( - objective=MultiObjective(objectives=[]) + objective=Objective(expression="branin_a, branin_b") ) moo_mixed = choose_generation_strategy_legacy( search_space=search_space, @@ -409,7 +409,7 @@ def test_choose_generation_strategy_legacy(self) -> None: num_initialization_trials=3, use_saasbo=True, optimization_config=MultiObjectiveOptimizationConfig( - objective=MultiObjective(objectives=[]) + objective=Objective(expression="branin_a, branin_b") ), ) self.assertEqual( diff --git a/ax/generation_strategy/tests/test_transition_criterion.py b/ax/generation_strategy/tests/test_transition_criterion.py index e7fd2b6428d..93563a08f0f 100644 --- a/ax/generation_strategy/tests/test_transition_criterion.py +++ b/ax/generation_strategy/tests/test_transition_criterion.py @@ -34,6 +34,7 @@ MaxTrialsAwaitingData, MinTrials, ) +from ax.metrics.branin import BraninMetric from ax.utils.common.constants import Keys from ax.utils.common.hash_utils import compute_lilo_input_hash from ax.utils.common.logger import get_logger @@ -440,6 +441,16 @@ def test_auto_with_should_skip_node(self) -> None: def test_is_single_objective_does_not_transition(self) -> None: exp = self.branin_experiment + exp.add_metric( + metric=BraninMetric( + name="branin_a", param_names=["x1", "x2"], lower_is_better=True + ) + ) + exp.add_metric( + metric=BraninMetric( + name="branin_b", param_names=["x1", "x2"], lower_is_better=True + ) + ) exp.optimization_config = get_branin_multi_objective_optimization_config() gs = GenerationStrategy( name="test", diff --git a/ax/global_stopping/strategies/improvement.py b/ax/global_stopping/strategies/improvement.py index 09760370a9c..6e7728586a3 100644 --- a/ax/global_stopping/strategies/improvement.py +++ b/ax/global_stopping/strategies/improvement.py @@ -16,7 +16,7 @@ MultiObjectiveOptimizationConfig, OptimizationConfig, ) -from ax.core.outcome_constraint import ObjectiveThreshold +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.trial import Trial from ax.core.types import ComparisonOp from ax.exceptions.core import AxError @@ -79,7 +79,7 @@ def __init__( self.window_size = window_size self.improvement_bar = improvement_bar self.hv_by_trial: dict[int, float] = {} - self._inferred_objective_thresholds: list[ObjectiveThreshold] | None = None + self._inferred_objective_thresholds: list[OutcomeConstraint] | None = None def __repr__(self) -> str: return super().__repr__() + ( @@ -93,7 +93,7 @@ def _should_stop_optimization( self, experiment: Experiment, trial_to_check: int | None = None, - objective_thresholds: list[ObjectiveThreshold] | None = None, + objective_thresholds: list[OutcomeConstraint] | None = None, ) -> tuple[bool, str]: """ Check if the objective has improved significantly in the past @@ -198,7 +198,7 @@ def _should_stop_moo( self, experiment: Experiment, trial_to_check: int, - objective_thresholds: list[ObjectiveThreshold], + objective_thresholds: list[OutcomeConstraint], ) -> tuple[bool, str]: """ This is the "should_stop_optimization" method of this class, specialized @@ -362,7 +362,7 @@ def constraint_satisfaction(trial: BaseTrial) -> bool: df = trial.lookup_data().df for constraint in outcome_constraints: bound = constraint.bound - metric_name = constraint.metric.name + metric_name = constraint.metric_names[0] metric_data = df.loc[df["metric_name"] == metric_name] mean, sem = metric_data.iloc[0][["mean", "sem"]] if sem > 0.0: diff --git a/ax/global_stopping/tests/test_strategies.py b/ax/global_stopping/tests/test_strategies.py index dcd8cfde0e9..e394e6adf38 100644 --- a/ax/global_stopping/tests/test_strategies.py +++ b/ax/global_stopping/tests/test_strategies.py @@ -221,13 +221,13 @@ def _create_multi_objective_experiment( ] objective_thresholds = [ ObjectiveThreshold( - metric=objectives[0].metric, + metric=Metric(name="m1"), bound=0.1, op=ComparisonOp.GEQ, relative=False, ), ObjectiveThreshold( - metric=objectives[1].metric, + metric=Metric(name="m2"), bound=0.2, op=ComparisonOp.GEQ, relative=False, @@ -245,6 +245,8 @@ def _create_multi_objective_experiment( optimization_config = MultiObjectiveOptimizationConfig( objective=MultiObjective(objectives), outcome_constraints=outcome_constraints, + # pyre-ignore[6]: ObjectiveThreshold is a subclass of OutcomeConstraint; + # list invariance prevents direct assignment. objective_thresholds=objective_thresholds, ) @@ -297,18 +299,16 @@ def test_multi_objective(self) -> None: gss = ImprovementGlobalStoppingStrategy( min_trials=3, window_size=3, improvement_bar=0.1 ) - objectives = assert_is_instance( - none_throws(exp.optimization_config).objective, MultiObjective - ).objectives + metric_names = list(none_throws(exp.optimization_config).objective.metric_names) custom_objective_thresholds = [ ObjectiveThreshold( - metric=objectives[0].metric, + metric=exp.get_metric(metric_names[0]), bound=-10, op=ComparisonOp.GEQ, relative=False, ), ObjectiveThreshold( - metric=objectives[1].metric, + metric=exp.get_metric(metric_names[1]), bound=-10, op=ComparisonOp.GEQ, relative=False, diff --git a/ax/metrics/tests/test_map_replay.py b/ax/metrics/tests/test_map_replay.py index 0089315f1d5..c2957bd0ae5 100644 --- a/ax/metrics/tests/test_map_replay.py +++ b/ax/metrics/tests/test_map_replay.py @@ -53,6 +53,7 @@ def test_map_replay(self) -> None: minimize=True, ) ), + tracking_metrics=[replay_metric], runner=SyntheticRunner(), ) for i in range(0, 2): @@ -114,6 +115,7 @@ def test_map_replay_non_uniform(self) -> None: minimize=True, ) ), + tracking_metrics=[replay_metric], runner=SyntheticRunner(), ) for i in range(0, 2): diff --git a/ax/orchestration/orchestrator.py b/ax/orchestration/orchestrator.py index a56357a3b4c..25203e51a98 100644 --- a/ax/orchestration/orchestrator.py +++ b/ax/orchestration/orchestrator.py @@ -1207,7 +1207,7 @@ def _error_if_status_quo_infeasible(self) -> None: if is_infeasible: constraint_descriptions = [ - f"{c.metric.name} {c.op.name} {c.bound}" + f"{c.metric_names[0]} {c.op.name} {c.bound}" for c in optimization_config.outcome_constraints ] error_msg = ( @@ -2097,7 +2097,7 @@ def _fetch_and_process_trials_data_results( optimization_config = self.experiment.optimization_config if ( optimization_config is not None - and metric_name in optimization_config.metrics.keys() + and metric_name in optimization_config.metric_names and not self.experiment.metrics[metric_name].is_recoverable_fetch_e( metric_fetch_e=metric_fetch_e ) diff --git a/ax/orchestration/tests/test_orchestrator.py b/ax/orchestration/tests/test_orchestrator.py index 8dbb7d79e50..0f7899e4298 100644 --- a/ax/orchestration/tests/test_orchestrator.py +++ b/ax/orchestration/tests/test_orchestrator.py @@ -31,6 +31,7 @@ from ax.core.objective import Objective from ax.core.observation import ObservationFeatures from ax.core.optimization_config import OptimizationConfig +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.runner import Runner from ax.core.utils import ( extract_pending_observations, @@ -322,7 +323,7 @@ def test_init_with_no_impl_with_runner(self) -> None: self.branin_experiment_no_impl_runner_or_metrics._optimization_config = None with self.assertRaisesRegex( UnsupportedError, - "`Orchestrator` requires that `experiment.metrics` not be None.", + "do not implement fetching logic", ): Orchestrator( experiment=self.branin_experiment_no_impl_runner_or_metrics, @@ -1547,7 +1548,7 @@ def test_get_best_trial(self) -> None: # results as expected, but only much deeper in the stack. # Python <3.14: "'branin_a' is not in list" # Python 3.14+: "list.index(x): x not in list" - with self.assertRaisesRegex(ValueError, "not in list"): + with self.assertRaisesRegex(KeyError, "not found"): orchestrator.get_pareto_optimal_parameters( optimization_config=get_branin_multi_objective_optimization_config( has_objective_thresholds=True @@ -1562,6 +1563,8 @@ def test_get_best_trial(self) -> None: self.assertEqual(params, just_params_unmodeled) def test_get_best_trial_moo(self) -> None: + for name in ["branin_a", "branin_b"]: + self.branin_experiment.add_tracking_metric(get_branin_metric(name=name)) self.branin_experiment.optimization_config = ( get_branin_multi_objective_optimization_config() ) @@ -2124,6 +2127,8 @@ def test_get_improvement_over_baseline(self) -> None: def test_get_improvement_over_baseline_robustness_not_implemented(self) -> None: """Test edge cases for get_improvement_over_baseline""" + for name in ["branin_a", "branin_b"]: + self.branin_experiment.add_tracking_metric(get_branin_metric(name=name)) self.branin_experiment.optimization_config = ( get_branin_multi_objective_optimization_config() ) @@ -2716,9 +2721,15 @@ def test_generate_candidates_updates_experiment_status(self) -> None: def test_generate_candidates_does_not_generate_if_missing_data(self) -> None: # GIVEN a orchestrator that can't fetch data + custom_metric = CustomTestMetric( + name="custom_test_metric", test_attribute="test" + ) + self.branin_experiment.add_tracking_metric(custom_metric) self.branin_experiment.optimization_config = OptimizationConfig( Objective( - CustomTestMetric(name="custom_test_metric", test_attribute="test"), + metric=CustomTestMetric( + name="custom_test_metric", test_attribute="test" + ), minimize=False, ) ) @@ -2858,7 +2869,9 @@ def test_terminate_if_status_quo_infeasible(self) -> None: experiment = get_branin_experiment( with_status_quo=True, with_absolute_constraint=True ) - experiment.optimization_config.outcome_constraints[0].bound = 100 + none_throws(experiment.optimization_config).outcome_constraints[0] = ( + OutcomeConstraint(expression="branin_e >= 100.0") + ) status_quo_trial = experiment.new_trial() status_quo_trial.add_arm(experiment.status_quo) @@ -2929,6 +2942,7 @@ def setUp(self) -> None: TestCase.setUp(self) self.branin_experiment = get_multi_type_experiment() self.branin_experiment.name = "branin_test_experiment" + self.branin_experiment.update_metric(BraninMetric("m1", ["x1", "x2"])) self.branin_experiment.optimization_config = OptimizationConfig( objective=Objective(metric=BraninMetric("m1", ["x1", "x2"]), minimize=True) ) @@ -2937,6 +2951,9 @@ def setUp(self) -> None: self.branin_experiment.update_runner(trial_type="type1", runner=self.runner) self.branin_timestamp_map_metric_experiment = get_multi_type_experiment() + self.branin_timestamp_map_metric_experiment.add_tracking_metric( + get_map_metric(name="branin_map") + ) self.branin_timestamp_map_metric_experiment.optimization_config = ( OptimizationConfig( objective=Objective( @@ -2951,7 +2968,7 @@ def setUp(self) -> None: self.branin_experiment_no_impl_runner_or_metrics = MultiTypeExperiment( search_space=get_branin_search_space(), optimization_config=OptimizationConfig( - Objective(Metric(name="branin"), minimize=True) + Objective(metric=Metric(name="branin"), minimize=True) ), default_trial_type="type1", default_runner=None, diff --git a/ax/plot/helper.py b/ax/plot/helper.py index 4eb9d239bff..2ab93d24064 100644 --- a/ax/plot/helper.py +++ b/ax/plot/helper.py @@ -795,7 +795,7 @@ def infer_is_relative( if model._optimization_config: constraints = none_throws(model._optimization_config).outcome_constraints constraint_relativity = { - constraint.metric.signature: constraint.relative + constraint.metric_names[0]: constraint.relative for constraint in constraints } for metric in metrics: diff --git a/ax/plot/pareto_frontier.py b/ax/plot/pareto_frontier.py index 1fd36a7d994..e24c1b4596b 100644 --- a/ax/plot/pareto_frontier.py +++ b/ax/plot/pareto_frontier.py @@ -584,9 +584,7 @@ def _validate_and_maybe_get_default_metric_names( multi_objective = assert_is_instance( none_throws(optimization_config).objective, MultiObjective ) - metric_names = tuple( - obj.metric.signature for obj in multi_objective.objectives - ) + metric_names = tuple(multi_objective.metric_names) else: raise UserInputError( "Inference of `metric_names` failed. Expected `MultiObjective` but " @@ -617,14 +615,14 @@ def _validate_experiment_and_maybe_get_objective_thresholds( if any( ot.relative for ot in objective_thresholds - if ot.metric.signature in metric_names + if ot.metric_names[0] in metric_names ): raise NotImplementedError( "Pareto plotting not supported for experiments with relative objective " "thresholds." ) constraint_metric_names = { - objective_threshold.metric.signature + objective_threshold.metric_names[0] for objective_threshold in objective_thresholds } missing_metric_names = set(metric_names) - set(constraint_metric_names) @@ -649,7 +647,7 @@ def _validate_and_maybe_get_default_reference_point( ) -> tuple[float, float] | None: if reference_point is None: reference_point = { - objective_threshold.metric.signature: objective_threshold.bound + objective_threshold.metric_names[0]: objective_threshold.bound for objective_threshold in objective_thresholds } missing_metric_names = set(metric_names) - set(reference_point) @@ -727,12 +725,9 @@ def _maybe_get_default_minimize_single_metric( and metric_name in optimization_config.objective.metric_names ): if optimization_config.is_moo_problem: - multi_objective = assert_is_instance( - optimization_config.objective, MultiObjective - ) - for objective in multi_objective.objectives: - if objective.metric.signature == metric_name: - return objective.minimize + for obj_metric_name, weight in optimization_config.objective.metric_weights: + if obj_metric_name == metric_name: + return weight < 0 else: return optimization_config.objective.minimize @@ -748,7 +743,7 @@ def _maybe_get_default_minimize_single_metric( f"{VALID_CONSTRAINT_OP_NAMES}. Got {invalid_constraint_op_names}.)" ) minimize = { - objective_threshold.metric.signature: objective_threshold.op.name == "LEQ" + objective_threshold.metric_names[0]: objective_threshold.op.name == "LEQ" for objective_threshold in objective_thresholds } minimize = minimize.get(metric_name) diff --git a/ax/plot/pareto_utils.py b/ax/plot/pareto_utils.py index a59150cc1e4..00930201995 100644 --- a/ax/plot/pareto_utils.py +++ b/ax/plot/pareto_utils.py @@ -200,10 +200,9 @@ def get_observed_pareto_frontiers( adapter = get_tensor_converter_adapter(experiment=experiment, data=data) pareto_observations = observed_pareto_frontier(adapter=adapter) # Convert to ParetoFrontierResults - objective_metric_names = { - metric.name - for metric in experiment.optimization_config.objective.metrics # pyre-ignore - } + objective_metric_names = set( + experiment.optimization_config.objective.metric_names # pyre-ignore + ) obj_metr_list = sorted(objective_metric_names) pfr_means = {name: [] for name in obj_metr_list} pfr_sems = {name: [] for name in obj_metr_list} @@ -221,8 +220,8 @@ def get_observed_pareto_frontiers( objective_thresholds = {} if experiment.optimization_config.objective_thresholds is not None: # pyre-ignore for objth in experiment.optimization_config.objective_thresholds: - rel_objth[objth.metric.signature] = objth.relative - objective_thresholds[objth.metric.signature] = objth.bound + rel_objth[objth.metric_names[0]] = objth.relative + objective_thresholds[objth.metric_names[0]] = objth.bound # Identify which metrics should be relativized if rel in [True, False]: @@ -492,10 +491,10 @@ def _validate_outcome_constraints( objective_metrics = [primary_objective.name, secondary_objective.name] if outcome_constraints is not None: for oc in outcome_constraints: - if oc.metric.signature in objective_metrics: + if oc.metric_names[0] in objective_metrics: raise ValueError( "Metric `{metric_name}` occurs in both outcome constraints " - "and objectives".format(metric_name=oc.metric.signature) + "and objectives".format(metric_name=oc.metric_names[0]) ) diff --git a/ax/plot/tests/test_pareto_utils.py b/ax/plot/tests/test_pareto_utils.py index fc5a0a8decd..739cd58ca57 100644 --- a/ax/plot/tests/test_pareto_utils.py +++ b/ax/plot/tests/test_pareto_utils.py @@ -10,6 +10,7 @@ import numpy as np from ax.adapter.registry import Generators from ax.core.data import Data +from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective from ax.core.optimization_config import MultiObjectiveOptimizationConfig from ax.core.outcome_constraint import ObjectiveThreshold @@ -76,7 +77,7 @@ def test_get_observed_pareto_frontiers(self) -> None: rels = [True, True, False] objective_thresholds = [ ObjectiveThreshold( - metric=objective.metric, + metric=Metric(name=objective.metric_names[0]), bound=bounds[i], relative=rels[i], op=ComparisonOp.LEQ, @@ -86,8 +87,24 @@ def test_get_observed_pareto_frontiers(self) -> None: objective = MultiObjective(objectives=objectives) optimization_config = MultiObjectiveOptimizationConfig( objective=objective, + # pyre-ignore[6]: ObjectiveThreshold is a subclass of OutcomeConstraint; + # list invariance prevents direct assignment. objective_thresholds=objective_thresholds, ) + + experiment.add_metric( + metric=BraninMetric( + name="m1", param_names=["x1", "x2"], lower_is_better=True + ) + ) + experiment.add_metric( + NegativeBraninMetric( + name="m2", param_names=["x1", "x2"], lower_is_better=True + ) + ) + experiment.add_metric( + BraninMetric(name="m3", param_names=["x1", "x2"], lower_is_better=True) + ) experiment.optimization_config = optimization_config experiment.trials[0].run() diff --git a/ax/runners/tests/test_map_replay.py b/ax/runners/tests/test_map_replay.py index e23dc36b481..5a9b6d8f2a1 100644 --- a/ax/runners/tests/test_map_replay.py +++ b/ax/runners/tests/test_map_replay.py @@ -45,6 +45,7 @@ def test_map_replay(self) -> None: ) ), runner=runner, + tracking_metrics=[metric], ) for _ in range(2): trial = experiment.new_trial() diff --git a/ax/runners/tests/test_torchx.py b/ax/runners/tests/test_torchx.py index f891c4b9eeb..cf369fbe590 100644 --- a/ax/runners/tests/test_torchx.py +++ b/ax/runners/tests/test_torchx.py @@ -57,10 +57,9 @@ def setUp(self) -> None: ] self._minimize = True + self._metric = TorchXMetric(name="booth_eval") self._objective = Objective( - metric=TorchXMetric( - name="booth_eval", - ), + metric=self._metric, minimize=self._minimize, ) @@ -85,6 +84,7 @@ def test_run_experiment_locally(self) -> None: runner=self._runner, is_test=True, properties={Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF: True}, + tracking_metrics=[self._metric], ) orchestrator = Orchestrator( @@ -116,6 +116,7 @@ def test_stop_trials(self) -> None: runner=self._runner, is_test=True, properties={Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF: True}, + tracking_metrics=[self._metric], ) orchestrator = Orchestrator( experiment=experiment, @@ -153,6 +154,7 @@ def test_run_experiment_locally_in_batches(self) -> None: runner=self._runner, is_test=True, properties={Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF: True}, + tracking_metrics=[self._metric], ) orchestrator = Orchestrator( @@ -185,6 +187,7 @@ def test_runner_no_batch_trials(self) -> None: optimization_config=OptimizationConfig(objective=self._objective), runner=self._runner, is_test=True, + tracking_metrics=[self._metric], ) with self.assertRaises(ValueError): diff --git a/ax/service/ax_client.py b/ax/service/ax_client.py index d7e3e3c1c44..b7fd4503f9d 100644 --- a/ax/service/ax_client.py +++ b/ax/service/ax_client.py @@ -27,13 +27,18 @@ from ax.core.experiment import Experiment from ax.core.generator_run import GeneratorRun from ax.core.multi_type_experiment import MultiTypeExperiment -from ax.core.objective import MultiObjective, Objective +from ax.core.objective import Objective from ax.core.observation import ObservationFeatures from ax.core.parameter import RangeParameter from ax.core.runner import Runner from ax.core.trial import Trial from ax.core.trial_status import TrialStatus -from ax.core.types import TEvaluationOutcome, TParameterization, TParamValue +from ax.core.types import ( + ComparisonOp, + TEvaluationOutcome, + TParameterization, + TParamValue, +) from ax.early_stopping.strategies import BaseEarlyStoppingStrategy from ax.early_stopping.utils import estimate_early_stopping_savings from ax.exceptions.constants import CHOLESKY_ERROR_ANNOTATION @@ -414,6 +419,32 @@ def set_optimization_config( metric_definitions=metric_definitions, ) if optimization_config: + # Build lower_is_better map from the optimization config so + # that auto-registered metrics carry the correct directionality. + lower_is_better_map: dict[str, bool] = {} + for mn, weight in optimization_config.objective.metric_weights: + lower_is_better_map[mn] = weight < 0 + for constraint in optimization_config.outcome_constraints: + for mn, _ in constraint.metric_weights: + if mn not in lower_is_better_map: + lower_is_better_map[mn] = constraint.op is ComparisonOp.LEQ + + # Auto-register metrics not yet on the experiment, using + # metric_definitions to preserve metric types and properties. + # For existing metrics, update lower_is_better to match the + # new optimization config. + for metric_name in optimization_config.metric_names: + if metric_name not in self.experiment.metrics: + metric = self._make_metric( + name=metric_name, + metric_definitions=metric_definitions, + lower_is_better=lower_is_better_map.get(metric_name), + ) + self.experiment.add_metric(metric) + elif metric_name in lower_is_better_map: + self.experiment.metrics[ + metric_name + ].lower_is_better = lower_is_better_map[metric_name] self.experiment.optimization_config = optimization_config self._save_experiment_to_db_if_possible( experiment=self.experiment, @@ -1009,7 +1040,7 @@ def get_optimization_trace( raise ValueError("Cannot generate plot as there are no trials.") objective = self.objective - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: raise UnsupportedError( "`get_optimization_trace` is not supported " "for multi-objective experiments" @@ -1085,7 +1116,7 @@ def get_contour_plot( ) if not metric_name: - if isinstance(self.objective, MultiObjective): + if self.objective.is_multi_objective: raise UnsupportedError( "`get_contour_plot` requires a `metric_name` " "for multi-objective experiments" @@ -1527,17 +1558,17 @@ def objective(self) -> Objective: def objective_name(self) -> str: """Returns the name of the objective in this optimization.""" objective = self.objective - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: raise UnsupportedError( "Multi-objective experiments contain multiple objectives" ) - return objective.metric.name + return objective.metric_names[0] @property def objective_names(self) -> list[str]: """Returns the name of the objective in this optimization.""" objective = self.objective - return [m.name for m in objective.metrics] + return list(objective.metric_names) @property def metric_definitions(self) -> dict[str, dict[str, Any]]: @@ -1781,7 +1812,7 @@ def _validate_all_required_metrics_present( trial_index=trial_index, metric_name_to_signature=metric_name_to_signature, ) - required_metrics = set(opt_config.metrics.keys()) + required_metrics = opt_config.metric_names provided_metrics = data.metric_names missing_metrics = required_metrics - provided_metrics return not missing_metrics diff --git a/ax/service/tests/test_ax_client.py b/ax/service/tests/test_ax_client.py index aaa48b3028d..b0afed53e05 100644 --- a/ax/service/tests/test_ax_client.py +++ b/ax/service/tests/test_ax_client.py @@ -383,28 +383,16 @@ def test_set_optimization_config_to_moo_with_constraints(self) -> None: ) objective = assert_is_instance(opt_config.objective, MultiObjective) self.assertEqual( - objective.objectives[0].metric.signature, + objective.metric_names[0], "foo", ) self.assertEqual( - objective.objectives[0].metric.name, - "foo", - ) - self.assertEqual( - objective.objectives[0].minimize, - True, - ) - self.assertEqual( - objective.objectives[1].metric.signature, + opt_config.objective.metric_names[1], "bar", ) self.assertEqual( - objective.objectives[1].metric.name, - "bar", - ) - self.assertEqual( - objective.objectives[1].minimize, - False, + [w < 0 for _, w in opt_config.objective.metric_weights], + [True, False], ) self.assertEqual( opt_config.objective_thresholds[0], @@ -452,7 +440,7 @@ def test_set_optimization_config_to_single_objective(self) -> None: ) opt_config = none_throws(ax_client.experiment.optimization_config) self.assertEqual( - opt_config.objective.metric.signature, + opt_config.objective.metric_names[0], "foo", ) self.assertEqual( @@ -891,10 +879,15 @@ def test_create_experiment(self) -> None: relative=False, ), ) - self.assertTrue(opt_config.objective.minimize) - self.assertDictEqual( - experiment._tracking_metrics, - {"test_tracking_metric": Metric(name="test_tracking_metric")}, + self.assertTrue( + none_throws( + none_throws(ax_client._experiment).optimization_config + ).objective.minimize + ) + self.assertEqual( + # pyre-fixme[16]: `Optional` has no attribute `tracking_metrics`. + [m.name for m in ax_client._experiment.tracking_metrics], + ["test_tracking_metric"], ) self.assertTrue(experiment.immutable_search_space_and_opt_config) self.assertTrue(ax_client.experiment.is_test) @@ -1106,31 +1099,35 @@ def test_create_experiment_with_metric_definitions(self) -> None: metric_definitions=metric_definitions, is_test=True, ) - opt_config = assert_is_instance( - ax_client.experiment.optimization_config, - MultiObjectiveOptimizationConfig, + # pyre-fixme[16]: `Optional` has no attribute `objective`. + objective = ax_client.experiment.optimization_config.objective + experiment = ax_client.experiment + self.assertEqual(objective.metric_names[0], "obj_m1") + self.assertEqual( + experiment.get_metric("obj_m1").properties, {"m1_opt": "m1_val"} ) - objective = assert_is_instance(opt_config.objective, MultiObjective) - objectives = objective.objectives - self.assertEqual(objectives[0].metric.signature, "obj_m1") - self.assertEqual(objectives[0].metric.name, "obj_m1") - self.assertEqual(objectives[0].metric.properties, {"m1_opt": "m1_val"}) - self.assertEqual(objectives[1].metric.signature, "obj_m2") - self.assertEqual(objectives[1].metric.name, "obj_m2") - self.assertEqual(objectives[1].metric.properties, {"m2_opt": "m2_val"}) - thresholds = opt_config.objective_thresholds - self.assertEqual(thresholds[0].metric.signature, "obj_m1") - self.assertEqual(thresholds[0].metric.name, "obj_m1") - - self.assertEqual(thresholds[0].metric.properties, {"m1_opt": "m1_val"}) - self.assertEqual(thresholds[1].metric.signature, "obj_m2") - self.assertEqual(thresholds[1].metric.name, "obj_m2") - - self.assertEqual(thresholds[1].metric.properties, {"m2_opt": "m2_val"}) + self.assertEqual(objective.metric_names[1], "obj_m2") + self.assertEqual( + experiment.get_metric("obj_m2").properties, {"m2_opt": "m2_val"} + ) + # pyre-fixme[16]: `Optional` has no attribute `objective_thresholds`. + thresholds = ax_client.experiment.optimization_config.objective_thresholds + self.assertEqual(thresholds[0].metric_names[0], "obj_m1") + + self.assertEqual( + experiment.get_metric("obj_m1").properties, {"m1_opt": "m1_val"} + ) + self.assertEqual(thresholds[1].metric_names[0], "obj_m2") + + self.assertEqual( + experiment.get_metric("obj_m2").properties, {"m2_opt": "m2_val"} + ) + opt_config = none_throws(experiment.optimization_config) outcome_constraints = opt_config.outcome_constraints - self.assertEqual(outcome_constraints[0].metric.signature, "const_m3") - self.assertEqual(outcome_constraints[0].metric.name, "const_m3") - self.assertEqual(outcome_constraints[0].metric.properties, {"m3_opt": "m3_val"}) + self.assertEqual(outcome_constraints[0].metric_names[0], "const_m3") + self.assertEqual( + experiment.get_metric("const_m3").properties, {"m3_opt": "m3_val"} + ) self.assertEqual( ax_client.experiment.tracking_metrics[0].signature, "tracking_m4" ) @@ -1211,29 +1208,42 @@ def test_set_optimization_config_with_metric_definitions(self) -> None: outcome_constraints=["const_m3 >= 3"], metric_definitions=metric_definitions, ) - opt_config = assert_is_instance( - ax_client.experiment.optimization_config, - MultiObjectiveOptimizationConfig, + # pyre-fixme[16]: `Optional` has no attribute `objective`. + objective = ax_client.experiment.optimization_config.objective + metric_names = objective.metric_names + self.assertEqual(metric_names[0], "obj_m1") + self.assertEqual( + ax_client.experiment.get_metric("obj_m1").properties, + {"m1_opt": "m1_val"}, + ) + self.assertEqual(metric_names[1], "obj_m2") + self.assertEqual( + ax_client.experiment.get_metric("obj_m2").properties, + {"m2_opt": "m2_val"}, + ) + # pyre-fixme[16]: `Optional` has no attribute `objective_thresholds`. + thresholds = ax_client.experiment.optimization_config.objective_thresholds + self.assertEqual(thresholds[0].metric_names[0], "obj_m1") + self.assertEqual( + ax_client.experiment.get_metric(thresholds[0].metric_names[0]).properties, + {"m1_opt": "m1_val"}, + ) + self.assertEqual(thresholds[1].metric_names[0], "obj_m2") + self.assertEqual( + ax_client.experiment.get_metric(thresholds[1].metric_names[0]).properties, + {"m2_opt": "m2_val"}, + ) + outcome_constraints = ( + # pyre-fixme[16]: `Optional` has no attribute `outcome_constraints`. + ax_client.experiment.optimization_config.outcome_constraints + ) + self.assertEqual(outcome_constraints[0].metric_names[0], "const_m3") + self.assertEqual( + ax_client.experiment.get_metric( + outcome_constraints[0].metric_names[0] + ).properties, + {"m3_opt": "m3_val"}, ) - objective = assert_is_instance(opt_config.objective, MultiObjective) - objectives = objective.objectives - self.assertEqual(objectives[0].metric.signature, "obj_m1") - self.assertEqual(objectives[0].metric.name, "obj_m1") - self.assertEqual(objectives[0].metric.properties, {"m1_opt": "m1_val"}) - self.assertEqual(objectives[1].metric.signature, "obj_m2") - self.assertEqual(objectives[1].metric.name, "obj_m2") - self.assertEqual(objectives[1].metric.properties, {"m2_opt": "m2_val"}) - thresholds = opt_config.objective_thresholds - self.assertEqual(thresholds[0].metric.signature, "obj_m1") - self.assertEqual(thresholds[0].metric.name, "obj_m1") - self.assertEqual(thresholds[0].metric.properties, {"m1_opt": "m1_val"}) - self.assertEqual(thresholds[1].metric.signature, "obj_m2") - self.assertEqual(thresholds[1].metric.name, "obj_m2") - self.assertEqual(thresholds[1].metric.properties, {"m2_opt": "m2_val"}) - outcome_constraints = opt_config.outcome_constraints - self.assertEqual(outcome_constraints[0].metric.signature, "const_m3") - self.assertEqual(outcome_constraints[0].metric.name, "const_m3") - self.assertEqual(outcome_constraints[0].metric.properties, {"m3_opt": "m3_val"}) self.assertEqual( ax_client.metric_definitions["obj_m1"]["properties"], metric_definitions["obj_m1"]["properties"], @@ -1644,19 +1654,15 @@ def test_create_moo_experiment(self) -> None: optimization_config.objective, MultiObjective ) self.assertEqual( - [m.name for m in multi_objective.metrics], + list(optimization_config.objective.metric_names), ["test_objective_1", "test_objective_2"], ) self.assertEqual( - [o.minimize for o in multi_objective.objectives], - [True, False], - ) - self.assertEqual( - [m.lower_is_better for m in multi_objective.metrics], + [w < 0 for _, w in multi_objective.metric_weights], [True, False], ) self.assertEqual( - [t.metric.name for t in optimization_config.objective_thresholds], + [t.metric_names[0] for t in optimization_config.objective_thresholds], ["test_objective_1", "test_objective_2"], ) self.assertEqual( @@ -1689,9 +1695,10 @@ def test_create_moo_experiment(self) -> None: relative=False, ), ) - self.assertDictEqual( - experiment._tracking_metrics, - {"test_tracking_metric": Metric(name="test_tracking_metric")}, + self.assertEqual( + # pyre-fixme[16]: `Optional` has no attribute `tracking_metrics`. + [m.name for m in ax_client._experiment.tracking_metrics], + ["test_tracking_metric"], ) self.assertTrue(experiment.immutable_search_space_and_opt_config) self.assertTrue(ax_client.experiment.is_test) diff --git a/ax/service/tests/test_best_point.py b/ax/service/tests/test_best_point.py index 33ab5338512..45ea14f67bd 100644 --- a/ax/service/tests/test_best_point.py +++ b/ax/service/tests/test_best_point.py @@ -26,6 +26,7 @@ MultiObjectiveOptimizationConfig, PreferenceOptimizationConfig, ) +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.parameter import ParameterType, RangeParameter from ax.core.search_space import SearchSpace from ax.core.trial import Trial @@ -39,6 +40,7 @@ ) from ax.service.utils.best_point_mixin import BestPointMixin from ax.utils.common.constants import Keys +from ax.utils.common.sympy import build_constraint_expression_str from ax.utils.common.testutils import TestCase from ax.utils.testing.core_stubs import ( get_experiment_with_batch_trial, @@ -58,7 +60,7 @@ def test_get_trace(self) -> None: # Same experiment with maximize via new optimization config. opt_conf = none_throws(exp.optimization_config).clone() - opt_conf.objective.minimize = False + opt_conf.objective = Objective(expression=opt_conf.objective.metric_names[0]) self.assertEqual(get_trace(exp, opt_conf), [11, 11, 11, 15, 15]) with self.subTest("Single objective with constraints"): @@ -113,8 +115,15 @@ def test_get_trace(self) -> None: # W/ relative constraints & status quo. exp.status_quo = Arm(parameters={"x": 0.5, "y": 0.5}, name="status_quo") - exp.optimization_config.outcome_constraints[0].bound = 1.0 - exp.optimization_config.outcome_constraints[0].relative = True + c = none_throws(exp.optimization_config).outcome_constraints[0] + none_throws(exp.optimization_config).outcome_constraints[0] = OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=1.0, + relative=True, + ) + ) # Fails if there's no data for status quo. with self.assertRaisesRegex(DataRequiredError, "relative constraint"): get_trace(exp) @@ -148,7 +157,15 @@ def test_get_trace(self) -> None: # test batch trial exp = get_experiment_with_batch_trial(with_status_quo=False) trial = exp.trials[0] - exp.optimization_config.outcome_constraints[0].relative = False + c = none_throws(exp.optimization_config).outcome_constraints[0] + none_throws(exp.optimization_config).outcome_constraints[0] = OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=c.bound, + relative=False, + ) + ) trial.mark_running(no_runner_required=True).mark_completed() df_dict = [] for i, arm in enumerate(trial.arms): @@ -417,7 +434,7 @@ def test_get_best_observed_value(self) -> None: self.assertEqual(get_best(exp), 5) # Same experiment with maximize via new optimization config. opt_conf = none_throws(exp.optimization_config).clone() - opt_conf.objective.minimize = False + opt_conf.objective = Objective(expression=opt_conf.objective.metric_names[0]) self.assertEqual(get_best(exp, opt_conf), 15) # Scalarized. @@ -629,10 +646,10 @@ def test_infer_reference_point_from_experiment(self) -> None: # be [-0.35, 0.35]. self.assertEqual(inferred_reference_point[0].op, ComparisonOp.LEQ) self.assertEqual(inferred_reference_point[0].bound, -0.35) - self.assertEqual(inferred_reference_point[0].metric.signature, "m1") + self.assertEqual(inferred_reference_point[0].metric_names[0], "m1") self.assertEqual(inferred_reference_point[1].op, ComparisonOp.GEQ) self.assertEqual(inferred_reference_point[1].bound, 0.35) - self.assertEqual(inferred_reference_point[1].metric.signature, "m2") + self.assertEqual(inferred_reference_point[1].metric_names[0], "m2") with mock.patch( "ax.service.utils.best_point.get_pareto_frontier_and_configs", @@ -659,7 +676,15 @@ def test_constrained_infer_reference_point_from_experiment(self) -> None: # TODO: Use experiment clone function once D50804778 is landed. experiment = copy.deepcopy(experiment) # Ensure that no observation is feasible. - experiment.optimization_config.outcome_constraints[0].bound = 1000.0 + oc_metric_name = ( + none_throws(experiment.optimization_config) + .outcome_constraints[0] + .metric_names[0] + ) + none_throws(experiment.optimization_config).outcome_constraints[0] = ( + OutcomeConstraint(expression=f"{oc_metric_name} >= 1000") + ) + experiments.append(experiment) for experiment in experiments: @@ -680,10 +705,10 @@ def test_constrained_infer_reference_point_from_experiment(self) -> None: # be [-0.35, 0.35]. self.assertEqual(inferred_reference_point[0].op, ComparisonOp.LEQ) self.assertEqual(inferred_reference_point[0].bound, -0.35) - self.assertEqual(inferred_reference_point[0].metric.signature, "m1") + self.assertEqual(inferred_reference_point[0].metric_names[0], "m1") self.assertEqual(inferred_reference_point[1].op, ComparisonOp.GEQ) self.assertEqual(inferred_reference_point[1].bound, 0.35) - self.assertEqual(inferred_reference_point[1].metric.signature, "m2") + self.assertEqual(inferred_reference_point[1].metric_names[0], "m2") def test_infer_reference_point_from_experiment_shuffled_metrics(self) -> None: # Generating an experiment with given data. @@ -760,10 +785,10 @@ def test_infer_reference_point_from_experiment_shuffled_metrics(self) -> None: self.assertEqual(inferred_reference_point[0].op, ComparisonOp.LEQ) self.assertEqual(inferred_reference_point[0].bound, -0.35) - self.assertEqual(inferred_reference_point[0].metric.signature, "m1") + self.assertEqual(inferred_reference_point[0].metric_names[0], "m1") self.assertEqual(inferred_reference_point[1].op, ComparisonOp.GEQ) self.assertEqual(inferred_reference_point[1].bound, 0.35) - self.assertEqual(inferred_reference_point[1].metric.signature, "m2") + self.assertEqual(inferred_reference_point[1].metric_names[0], "m2") def test_get_tensor_converter_adapter(self) -> None: # Test that it can convert experiments with different number of observations. diff --git a/ax/service/tests/test_best_point_utils.py b/ax/service/tests/test_best_point_utils.py index 1a14581de99..f5e9309ab49 100644 --- a/ax/service/tests/test_best_point_utils.py +++ b/ax/service/tests/test_best_point_utils.py @@ -14,7 +14,6 @@ import numpy as np import pandas as pd -import torch from ax.adapter.cross_validation import AssessModelFitResult from ax.adapter.registry import Generators from ax.adapter.torch import TorchAdapter @@ -47,6 +46,7 @@ ) from ax.service.utils.best_point_utils import select_baseline_name_default_first_trial from ax.service.utils.instantiation import ObjectiveProperties +from ax.utils.common.sympy import build_constraint_expression_str from ax.utils.common.testutils import TestCase from ax.utils.testing.core_stubs import ( get_branin_experiment, @@ -106,8 +106,15 @@ def test_get_values_of_outcomes_single_or_scalarized_objective(self) -> None: df_wide=df, objective=objective ) self.assertEqual(len(value), 2) - self.assertAlmostEqual(value[0], 0.4 * m1_value + 0.6 * m2_value) - self.assertEqual(value[1], worst) + # ScalarizedObjective bakes the minimize sign into weights, + # so the expression is always "maximize". When minimize=True, + # weights are negated, and infeasible entries get -inf. + sign = -1 if minimize else 1 + scalarized_worst = float("-inf") + self.assertAlmostEqual( + value[0], sign * (0.4 * m1_value + 0.6 * m2_value) + ) + self.assertEqual(value[1], scalarized_worst) objective = ScalarizedObjective( metrics=metrics, weights=[0.4, -0.6], minimize=minimize @@ -119,8 +126,8 @@ def test_get_values_of_outcomes_single_or_scalarized_objective(self) -> None: df_wide=df, objective=objective ) self.assertEqual(len(value), 2) - self.assertEqual(value[0], 0.4 * m1_value - 0.6 * m2_value) - self.assertEqual(value[1], worst) + self.assertEqual(value[0], sign * (0.4 * m1_value - 0.6 * m2_value)) + self.assertEqual(value[1], scalarized_worst) def test_get_hypervolume_trace_of_outcomes_multi_objective(self) -> None: objective = MultiObjective( @@ -219,7 +226,10 @@ def test_get_trace_by_arm_pull_from_data(self) -> None: objective=objective, outcome_constraints=[ OutcomeConstraint( - Metric("m2"), op=ComparisonOp.GEQ, bound=0.0, relative=False + metric=Metric("m2"), + op=ComparisonOp.GEQ, + bound=0.0, + relative=False, ) ], ) @@ -513,17 +523,6 @@ def test_best_raw_objective_point(self) -> None: best_metrics = none_throws(best_prediction)[0] self.assertDictEqual(best_metrics, {"m1": 3.0, "m2": 4.0}) - # Tensor bounds are accepted. - constraint = none_throws(exp.optimization_config).all_constraints[0] - # pyre-fixme[8]: Attribute `bound` declared in class `OutcomeConstraint` - # has type `float` but is used as type `Tensor`. - constraint.bound = torch.tensor(constraint.bound) - _, __, best_prediction = none_throws( - get_best_by_raw_objective_with_trial_index(exp) - ) - best_metrics = none_throws(best_prediction)[0] - self.assertDictEqual(best_metrics, {"m1": 3.0, "m2": 4.0}) - def test_best_raw_objective_point_unsatisfiable(self) -> None: exp = get_branin_experiment() trial = exp.new_trial( @@ -587,8 +586,15 @@ def test_best_raw_objective_point_unsatisfiable_relative(self) -> None: # Create altered optimization config with unsatisfiable relative constraint. opt_conf = none_throws(exp.optimization_config).clone() - opt_conf.outcome_constraints[0].relative = True - opt_conf.outcome_constraints[0].bound = 9999 + old_c = opt_conf.outcome_constraints[0] + opt_conf.outcome_constraints[0] = OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=old_c.metric_weights, + op=">=" if old_c.op == ComparisonOp.GEQ else "<=", + bound=9999, + relative=True, + ) + ) with self.assertRaisesRegex( DataRequiredError, @@ -664,7 +670,16 @@ def test_derelativize_opt_config(self) -> None: optimization_config=input_optimization_config, experiment=exp, ) - self.assertEqual(input_optimization_config, returned_opt_config) + # clone() may return base Objective instead of MultiObjective, + # so compare by expression rather than type-strict equality. + self.assertEqual( + input_optimization_config.objective.expression, + returned_opt_config.objective.expression, + ) + self.assertEqual( + input_optimization_config.outcome_constraints, + returned_opt_config.outcome_constraints, + ) with self.subTest("Provided opt config overrides experiment opt config"): opt_config = input_optimization_config.clone_with_args( @@ -674,12 +689,35 @@ def test_derelativize_opt_config(self) -> None: optimization_config=opt_config, experiment=exp, ) - self.assertEqual(opt_config, returned_opt_config) - self.assertNotEqual(opt_config, exp.optimization_config) + self.assertEqual( + opt_config.objective.expression, + returned_opt_config.objective.expression, + ) + self.assertEqual( + opt_config.outcome_constraints, + returned_opt_config.outcome_constraints, + ) + self.assertNotEqual( + opt_config.outcome_constraints, + none_throws(exp.optimization_config).outcome_constraints, + ) # Add relative constraints. - for constraint in input_optimization_config.all_constraints: - constraint.relative = True + def _to_relative(c: OutcomeConstraint) -> OutcomeConstraint: + return OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=c.bound, + relative=True, + ) + ) + + for i, c in enumerate(input_optimization_config.outcome_constraints): + input_optimization_config.outcome_constraints[i] = _to_relative(c) + if isinstance(input_optimization_config, MultiObjectiveOptimizationConfig): + for i, c in enumerate(input_optimization_config.objective_thresholds): + input_optimization_config.objective_thresholds[i] = _to_relative(c) # Check errors. with ( @@ -703,12 +741,8 @@ def test_derelativize_opt_config(self) -> None: ), MultiObjectiveOptimizationConfig, ) - relativized_obj = assert_is_instance( - relativized_opt_config.objective, MultiObjective - ) - input_obj = assert_is_instance( - input_optimization_config.objective, MultiObjective - ) + relativized_obj = relativized_opt_config.objective + input_obj = input_optimization_config.objective status_quo_df = exp.lookup_data(trial_indices={status_quo_trial_index}).df # This is not a real test of `derelativize_opt_config` but rather # making sure the values on the experiment have't drifted @@ -722,9 +756,11 @@ def test_derelativize_opt_config(self) -> None: for i in range(2): # Relativization is with respect to the absolute value of the # status quo value + relativized_weights = [w for _, w in relativized_obj.metric_weights] + input_weights = [w for _, w in input_obj.metric_weights] self.assertEqual( - relativized_obj.objectives[i].minimize, - input_obj.objectives[i].minimize, + relativized_weights[i] < 0, + input_weights[i] < 0, ) self.assertEqual( relativized_opt_config.objective_thresholds[i].bound, @@ -796,11 +832,20 @@ def _test_is_row_feasible( feasible_series, expected_series, check_names=False ) - exp.optimization_config.outcome_constraints[0].relative = True + c = none_throws(exp.optimization_config).outcome_constraints[0] + opt_config = none_throws(exp.optimization_config) + opt_config.outcome_constraints[0] = OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=c.bound, + relative=True, + ) + ) relative_constraint_warning = ( "WARNING:ax.service.utils.best_point:Determining trial feasibility only " "supported with a derelativized OptimizationConfig, but found the " - "following relative constraints: [OutcomeConstraint(m3 >= 0.0%)]. " + "following relative constraints: [OutcomeConstraint(m3 >= baseline)]. " f"Returning {undetermined_value} as the feasibility." ) with self.assertLogs(logger=best_point_logger, level="WARN") as lg: @@ -825,8 +870,26 @@ def _test_is_row_feasible( feasible_series, expected_series, check_names=False ) exp._status_quo = exp.trials[0].arms[0] - for constraint in none_throws(exp.optimization_config).all_constraints: - constraint.relative = True + opt_config = none_throws(exp.optimization_config) + for i, c in enumerate(opt_config.outcome_constraints): + opt_config.outcome_constraints[i] = OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=c.bound, + relative=True, + ) + ) + if isinstance(opt_config, MultiObjectiveOptimizationConfig): + for i, c in enumerate(opt_config.objective_thresholds): + opt_config.objective_thresholds[i] = OutcomeConstraint( + expression=build_constraint_expression_str( + metric_weights=c.metric_weights, + op=">=" if c.op == ComparisonOp.GEQ else "<=", + bound=c.bound, + relative=True, + ) + ) optimization_config = derelativize_opt_config( optimization_config=none_throws(exp.optimization_config), experiment=exp, diff --git a/ax/service/tests/test_instantiation_utils.py b/ax/service/tests/test_instantiation_utils.py index 800be9ab2e5..b8c1aaa0f4f 100644 --- a/ax/service/tests/test_instantiation_utils.py +++ b/ax/service/tests/test_instantiation_utils.py @@ -9,7 +9,6 @@ from collections.abc import Sequence from typing import Any -from ax.core.metric import Metric from ax.core.optimization_config import MultiObjectiveOptimizationConfig from ax.core.parameter import ( ChoiceParameter, @@ -197,16 +196,16 @@ def test_add_tracking_metrics(self) -> None: parameters=[{"name": "x", "type": "range", "bounds": [0, 1]}], tracking_metric_names=None, ) - self.assertDictEqual(experiment._tracking_metrics, {}) + self.assertEqual(experiment.tracking_metrics, []) metrics_names = ["metric_1", "metric_2"] experiment = InstantiationBase.make_experiment( parameters=[{"name": "x", "type": "range", "bounds": [0, 1]}], tracking_metric_names=metrics_names, ) - self.assertDictEqual( - experiment._tracking_metrics, - {metric_name: Metric(name=metric_name) for metric_name in metrics_names}, + self.assertCountEqual( + [m.name for m in experiment.tracking_metrics], + metrics_names, ) def test_make_objectives(self) -> None: @@ -219,10 +218,10 @@ def test_make_objectives(self) -> None: objectives = InstantiationBase.make_objectives( {"branin": "minimize", "currin": "maximize"} ) - branin_metric = [o.minimize for o in objectives if o.metric.name == "branin"] - self.assertTrue(branin_metric[0]) - currin_metric = [o.minimize for o in objectives if o.metric.name == "currin"] - self.assertFalse(currin_metric[0]) + branin_obj = [o for o in objectives if o.metric_names[0] == "branin"] + self.assertTrue(branin_obj[0].minimize) + currin_obj = [o for o in objectives if o.metric_names[0] == "currin"] + self.assertFalse(currin_obj[0].minimize) def test_make_optimization_config(self) -> None: objectives = {"branin": "minimize", "currin": "maximize"} @@ -243,7 +242,7 @@ def test_make_optimization_config(self) -> None: outcome_constraints=[], status_quo_defined=False, ) - self.assertEqual(len(multi_optimization_config.objective.metrics), 2) + self.assertEqual(len(multi_optimization_config.objective.metric_names), 2) self.assertEqual( len( assert_is_instance( @@ -260,7 +259,7 @@ def test_make_optimization_config(self) -> None: outcome_constraints=[], status_quo_defined=False, ) - self.assertEqual(len(multi_optimization_config.objective.metrics), 2) + self.assertEqual(len(multi_optimization_config.objective.metric_names), 2) self.assertEqual( len( assert_is_instance( @@ -279,9 +278,11 @@ def test_make_optimization_config(self) -> None: outcome_constraints=[], status_quo_defined=False, ) - self.assertEqual(single_optimization_config.objective.metric.name, "branin") self.assertEqual( - single_optimization_config.objective.metric.signature, "branin" + single_optimization_config.objective.metric_names[0], "branin" + ) + self.assertEqual( + single_optimization_config.objective.metric_names[0], "branin" ) def test_single_valued_choice_to_fixed_param_conversion(self) -> None: diff --git a/ax/service/tests/test_report_utils.py b/ax/service/tests/test_report_utils.py index 4eab606a5ef..4426c7b520f 100644 --- a/ax/service/tests/test_report_utils.py +++ b/ax/service/tests/test_report_utils.py @@ -395,8 +395,11 @@ def test_get_standard_plots(self) -> None: @mock_botorch_optimize def test_get_standard_plots_moo(self) -> None: exp = get_branin_experiment_with_multi_objective(with_batch=True) - exp.optimization_config.objective.objectives[0].minimize = False - exp.optimization_config.objective.objectives[1].minimize = True + # Flip the minimize direction for the multi-objective: + # first objective to maximize, second to minimize + obj = none_throws(exp.optimization_config).objective + names = obj.metric_names + obj._expression_str = f"{names[0]}, -{names[1]}" assert_is_instance( exp.optimization_config, MultiObjectiveOptimizationConfig )._objective_thresholds = [ @@ -436,29 +439,35 @@ def _test_get_standard_plots_moo_relative_constraints( self, trial_is_complete: bool ) -> None: exp = get_branin_experiment_with_multi_objective(with_batch=True) - first_obj = assert_is_instance( - none_throws(exp.optimization_config).objective, MultiObjective - ).objectives[0] - first_obj.minimize = False - first_obj.metric.lower_is_better = False + obj = none_throws(exp.optimization_config).objective + # Flip the first objective to maximize (positive weight = maximize) + names = obj.metric_names + # Create a new Objective rather than mutating _expression_str to + # avoid stale _parsed cached_property. + none_throws(exp.optimization_config)._objective = Objective( + expression=f"{names[0]}, -{names[1]}" + ) + exp.get_metric(names[0]).lower_is_better = False assert_is_instance( exp.optimization_config, MultiObjectiveOptimizationConfig )._objective_thresholds = [ ObjectiveThreshold( - metric=exp.metrics["branin_a"], op=ComparisonOp.GEQ, bound=-100.0 + metric=exp.metrics["branin_a"], + op=ComparisonOp.GEQ, + bound=-100.0, + relative=False, ), ObjectiveThreshold( - metric=exp.metrics["branin_b"], op=ComparisonOp.LEQ, bound=100.0 + metric=exp.metrics["branin_b"], + op=ComparisonOp.LEQ, + bound=100.0, + relative=False, ), ] exp.trials[0].run() if trial_is_complete: exp.trials[0].mark_completed() - for ot in assert_is_instance( - exp.optimization_config, MultiObjectiveOptimizationConfig - )._objective_thresholds: - ot.relative = False plots = get_standard_plots( experiment=exp, model=Generators.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data()), @@ -475,8 +484,13 @@ def test_get_standard_plots_moo_relative_constraints(self) -> None: @mock_botorch_optimize def test_get_standard_plots_moo_no_objective_thresholds(self) -> None: exp = get_branin_experiment_with_multi_objective(with_batch=True) - exp.optimization_config.objective.objectives[0].minimize = False - exp.optimization_config.objective.objectives[1].minimize = True + # Flip the minimize direction for the multi-objective: + # first objective to maximize, second to minimize + obj = none_throws(exp.optimization_config).objective + names = obj.metric_names + none_throws(exp.optimization_config)._objective = Objective( + expression=f"{names[0]}, -{names[1]}" + ) exp.trials[0].run() plots = get_standard_plots( experiment=exp, @@ -641,7 +655,6 @@ def test_warn_if_unpredictable_metrics(self) -> None: self.assertIsNone(msg) # Test with no optimization config. - exp._tracking_metrics = exp.metrics exp._optimization_config = None msg = warn_if_unpredictable_metrics( experiment=exp, @@ -701,7 +714,7 @@ def test_maybe_extract_baseline_comparison_values_metric_missing_soo( arm_names = list(exp.arms_by_name.keys()) # Use a different metric name in optimization config that doesn't exist in data - exp._optimization_config.objective._metric = Metric(name="nonexistent_metric") + exp._optimization_config.objective._expression_str = "nonexistent_metric" result = maybe_extract_baseline_comparison_values( experiment=exp, @@ -720,7 +733,8 @@ def test_maybe_extract_baseline_comparison_values_metric_missing_moo(self) -> No # Replace one objective metric with a nonexistent one moo = none_throws(exp.optimization_config).objective - moo.objectives[0]._metric = Metric(name="nonexistent_metric") + names = list(moo.metric_names) + moo._expression_str = f"-nonexistent_metric, -{names[1]}" result = maybe_extract_baseline_comparison_values( experiment=exp, diff --git a/ax/service/utils/best_point.py b/ax/service/utils/best_point.py index be96ef6d81d..0a9fb6c90cc 100644 --- a/ax/service/utils/best_point.py +++ b/ax/service/utils/best_point.py @@ -35,14 +35,14 @@ from ax.core.data import Data from ax.core.experiment import Experiment from ax.core.generator_run import GeneratorRun -from ax.core.objective import MultiObjective, Objective, ScalarizedObjective +from ax.core.objective import Objective from ax.core.observation import ObservationFeatures from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, OptimizationConfig, PreferenceOptimizationConfig, ) -from ax.core.outcome_constraint import ObjectiveThreshold, OutcomeConstraint +from ax.core.outcome_constraint import _op_to_str, OutcomeConstraint from ax.core.trial import Trial from ax.core.types import ComparisonOp, TModelPredictArm, TParameterization from ax.exceptions.core import DataRequiredError, UnsupportedError, UserInputError @@ -172,8 +172,10 @@ def get_best_raw_objective_point_with_trial_index( experiment=experiment, ) - maximize = isinstance(optimization_config.objective, MultiObjective) or ( - not optimization_config.objective.minimize + maximize = ( + isinstance(optimization_config, MultiObjectiveOptimizationConfig) + or optimization_config.objective.is_scalarized_objective + or not optimization_config.objective.minimize ) best_row_idx = ( value_by_arm_pull["value"].idxmax() @@ -303,12 +305,7 @@ def get_best_parameters_from_model_predictions_with_trial_index( # For ScalarizedObjective, check model fit for all component metrics # For regular Objective, check model fit for the single objective metric - if isinstance(optimization_config.objective, ScalarizedObjective): - objective_metric_names = [ - metric.name for metric in optimization_config.objective.metrics - ] - else: - objective_metric_names = [optimization_config.objective.metric.name] + objective_metric_names = list(optimization_config.objective.metric_names) # If model fit is bad for any objective metric, use raw results bad_fit_objective_metrics = [ @@ -481,9 +478,11 @@ def get_pareto_optimal_parameters( ) adapter = assert_is_instance(adapter, TorchAdapter) - objective_thresholds_override = None + objective_thresholds_override: list[OutcomeConstraint] | None = None # If objective thresholds are not specified in optimization config, infer them. if not moo_optimization_config.objective_thresholds: + # pyre-ignore[9]: ObjectiveThreshold is a subclass of OutcomeConstraint; + # list invariance prevents direct assignment. objective_thresholds_override = adapter.infer_objective_thresholds( search_space=experiment.search_space, optimization_config=optimization_config, @@ -578,7 +577,7 @@ def compute_feasibility_per_constraint( upper_bound: pd.Series = upper_bound, name: pd.Series = name, ) -> pd.Series: - name_match_mask = name == oc.metric.name + name_match_mask = name == oc.metric_names[0] # Return True if metrics are different, or whether the confidence # interval is entirely not within the bound if oc.op == ComparisonOp.GEQ: @@ -594,7 +593,7 @@ def compute_feasibility_per_constraint( # conclusively. has_missing_metric_mask = pd.Series([False] * len(df), index=df.index) constrained_metric_names = { - oc.metric.name for oc in optimization_config.all_constraints + oc.metric_names[0] for oc in optimization_config.all_constraints } for arm_name, arm_group in df.groupby("arm_name"): metrics_for_arm = set(arm_group["metric_name"].unique()) @@ -679,18 +678,28 @@ def get_values_of_outcomes_single_or_scalarized_objective( ... ) np.array([2.0, inf]) """ - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: raise ValueError( "MultiObjective is not supported. Use " "`get_hypervolume_of_outcomes_multi_objective`." ) - if isinstance(objective, ScalarizedObjective): - value = df_wide[objective.metric_names].dot(objective.weights).to_numpy() + if objective.is_scalarized_objective: + value = ( + df_wide[objective.metric_names] + .dot([w for _, w in objective.metric_weights]) + .to_numpy() + ) else: - value = df_wide[objective.metric.name].to_numpy() + value = df_wide[objective.metric_names[0]].to_numpy() value = value.astype(np.float64) infeasible_idx = np.where(~df_wide["feasible"])[0] - value[infeasible_idx] = float("inf") if objective.minimize else float("-inf") + # For scalarized objectives the expression is always formulated as + # "maximize", so infeasible values should be -inf. + if objective.is_scalarized_objective: + minimize = False + else: + minimize = objective.minimize + value[infeasible_idx] = float("inf") if minimize else float("-inf") return value @@ -795,20 +804,19 @@ def get_hypervolume_trace_of_outcomes_multi_objective( ... ) [0.0, 2.0, 0.0, 9.0] """ - objective = assert_is_instance(optimization_config.objective, MultiObjective) - for obj in objective.objectives: - if obj.minimize: - df_wide[obj.metric.name] *= -1 + objective = optimization_config.objective + for metric_name, weight in objective.metric_weights: + if weight < 0: # minimize + df_wide[metric_name] *= -1 objective_thresholds = [] objective_thresholds_dict = { - threshold.metric.name: threshold + threshold.metric_names[0]: threshold for threshold in optimization_config.objective_thresholds } # First pass: collect explicit thresholds, mark missing ones with NaN. needs_inference = False - for obj in objective.objectives: - metric_name = obj.metric.name + for metric_name, weight in objective.metric_weights: if metric_name in objective_thresholds_dict: threshold = objective_thresholds_dict[metric_name] if threshold.relative: @@ -820,7 +828,7 @@ def get_hypervolume_trace_of_outcomes_multi_objective( # Explicit thresholds are in the original metric space, so negate # for minimization objectives to match the negated data. bound = threshold.bound - objective_thresholds.append(-bound if obj.minimize else bound) + objective_thresholds.append(-bound if weight < 0 else bound) else: needs_inference = True objective_thresholds.append(float("nan")) @@ -980,7 +988,7 @@ def _compute_trace_values( """ objective = optimization_config.objective # MOO and *not* ScalarizedObjective - if isinstance(objective, MultiObjective): + if objective.is_multi_objective: maximize = True optimization_config = assert_is_instance( optimization_config, MultiObjectiveOptimizationConfig @@ -993,7 +1001,11 @@ def _compute_trace_values( ) ) else: - maximize = not objective.minimize + if objective.is_scalarized_objective: + # Scalarized expressions are always in maximization convention + maximize = True + else: + maximize = not objective.minimize values = pd.Series( get_values_of_outcomes_single_or_scalarized_objective( df_wide=df_wide, objective=objective @@ -1028,7 +1040,7 @@ def _pivot_data_with_feasibility( where "feasible" indicates whether the observation satisfies all constraints. """ # Get the metrics we need - metric_names = list(optimization_config.metrics.keys()) + metric_names = list(optimization_config.metric_names) mask = df["metric_name"].isin(metric_names) # Transform to wide format with metric columns @@ -1191,7 +1203,6 @@ def get_trace_by_arm_pull_from_data( optimization_config=optimization_config, use_cumulative_best=use_cumulative_best, ) - return df_wide[["trial_index", "arm_name", "value"]] @@ -1202,7 +1213,6 @@ def get_trace( ) -> list[float]: """ Compute the optimization trace at each iteration. - Given an experiment and an optimization config, compute the performance at each iteration. For multi-objective, the performance is computed as the hypervolume. For single objective, the performance is computed as @@ -1243,7 +1253,7 @@ def get_trace( # Get the names of the metrics in optimization config. metric_names = set(optimization_config.objective.metric_names) for cons in optimization_config.outcome_constraints: - metric_names.update({cons.metric.name}) + metric_names.update(set(cons.metric_names)) metric_names = list(metric_names) # Don't compute results for status quo data (for compatibility with legacy behavior) @@ -1279,7 +1289,11 @@ def get_trace( ) # Aggregate by trial, then. compute cumulative best objective = optimization_config.objective - maximize = isinstance(objective, MultiObjective) or not objective.minimize + maximize = ( + isinstance(optimization_config, MultiObjectiveOptimizationConfig) + or objective.is_scalarized_objective + or not objective.minimize + ) cumulative_value = _aggregate_and_cumulate_trace( df=value_by_arm_pull, by=["trial_index"], @@ -1345,7 +1359,7 @@ def get_tensor_converter_adapter( def infer_reference_point_from_experiment( experiment: Experiment, data: Data -) -> list[ObjectiveThreshold]: +) -> list[OutcomeConstraint]: """This functions is a wrapper around ``infer_reference_point`` to find the nadir point from the pareto front of an experiment. Aside from converting experiment to tensors, this wrapper transforms back and forth the objectives of the experiment @@ -1389,32 +1403,37 @@ def infer_reference_point_from_experiment( inferred_rp = deepcopy(opt_config.objective_thresholds) else: inferred_rp = [] - for objective in assert_is_instance( - opt_config.objective, MultiObjective - ).objectives: - ot = ObjectiveThreshold( - metric=objective.metric, - bound=0.0, # dummy value - op=ComparisonOp.LEQ if objective.minimize else ComparisonOp.GEQ, + for metric_name, weight in opt_config.objective.metric_weights: + minimize = weight < 0 + op_str = "<=" if minimize else ">=" + ot = OutcomeConstraint( + expression=f"{metric_name} {op_str} 0.0", relative=False, ) inferred_rp.append(ot) + + processed_rp = [] for ot in inferred_rp: # In the following, we find the index of the objective in # `objective_orders`. If there is an objective that does not exist # in `obs_data`, a ValueError is raised. try: - objective_index = objective_orders.index(ot.metric.signature) + objective_index = objective_orders.index(ot.metric_names[0]) except ValueError: raise ValueError( - f"Metric {ot.metric.signature} does not exist in `obs_data`." + f"Metric {ot.metric_names[0]} does not exist in `obs_data`." ) + # Note: Sympy interprets oo as infinity if ot.op == ComparisonOp.LEQ: - ot.bound = np.inf + processed_rp.append( + OutcomeConstraint(expression=f"{ot.metric_names[0]} <= oo") + ) multiplier[objective_index] = -1 else: - ot.bound = -np.inf + processed_rp.append( + OutcomeConstraint(expression=f"{ot.metric_names[0]} >= -oo") + ) multiplier[objective_index] = 1 # Finding the pareto frontier @@ -1422,7 +1441,7 @@ def infer_reference_point_from_experiment( adapter=adapter, observation_features=obs_feats, observation_data=obs_data, - objective_thresholds=inferred_rp, + objective_thresholds=processed_rp, use_model_predictions=False, ) if len(frontier_observations) == 0: @@ -1445,7 +1464,7 @@ def infer_reference_point_from_experiment( adapter=adapter, observation_features=obs_feats, observation_data=obs_data, - objective_thresholds=inferred_rp, + objective_thresholds=processed_rp, use_model_predictions=False, ) # restoring constraints @@ -1482,16 +1501,17 @@ def infer_reference_point_from_experiment( x for (i, x) in enumerate(objective_orders) if multiplier[i] != 0 ] - for obj_threshold in inferred_rp: - obj_threshold.bound = rp[ - objective_orders_reduced.index(obj_threshold.metric.signature) - ].item() - return inferred_rp + return [ + OutcomeConstraint( + expression=f"{obj_threshold.metric_names[0]} {_op_to_str(obj_threshold.op)} {rp[objective_orders_reduced.index(obj_threshold.metric_names[0])].item()}" # noqa: E501 + ) + for obj_threshold in inferred_rp + ] def _get_objective_thresholds( optimization_config: MultiObjectiveOptimizationConfig, -) -> list[ObjectiveThreshold]: +) -> list[OutcomeConstraint]: """Get objective thresholds for an optimization config. This will return objective thresholds with dummy values if there are @@ -1506,13 +1526,11 @@ def _get_objective_thresholds( if optimization_config.objective_thresholds is not None: return deepcopy(optimization_config.objective_thresholds) objective_thresholds = [] - for objective in assert_is_instance( - optimization_config.objective, MultiObjective - ).objectives: - ot = ObjectiveThreshold( - metric=objective.metric, - bound=0.0, # dummy value - op=ComparisonOp.LEQ if objective.minimize else ComparisonOp.GEQ, + for metric_name, weight in optimization_config.objective.metric_weights: + minimize = weight < 0 + op_str = "<=" if minimize else ">=" + ot = OutcomeConstraint( + expression=f"{metric_name} {op_str} 0.0", relative=False, ) objective_thresholds.append(ot) diff --git a/ax/service/utils/best_point_mixin.py b/ax/service/utils/best_point_mixin.py index 2f9fef1b028..66d1a4314c5 100644 --- a/ax/service/utils/best_point_mixin.py +++ b/ax/service/utils/best_point_mixin.py @@ -306,8 +306,8 @@ def _get_best_observed_value( objective = optimization_config.objective if isinstance(objective, ScalarizedObjective): value = 0 - for metric, weight in objective.metric_weights: - value += means[metric.name] * weight + for metric_name, weight in objective.metric_weights: + value += means[metric_name] * weight return value else: name = objective.metric_names[0] @@ -385,7 +385,7 @@ def _get_trace_by_progression( optimization_config = optimization_config or none_throws( experiment.optimization_config ) - objective = optimization_config.objective.metric.name + objective = optimization_config.objective.metric_names[0] minimize = optimization_config.objective.minimize map_data = experiment.lookup_data() if not map_data.has_step_column: @@ -468,7 +468,7 @@ def get_improvement_over_baseline( if not optimization_config: raise ValueError("No optimization config found.") - objective_metric_name = optimization_config.objective.metric.name + objective_metric_name = optimization_config.objective.metric_names[0] # get the baseline trial data = experiment.lookup_data().df diff --git a/ax/service/utils/instantiation.py b/ax/service/utils/instantiation.py index 8fe8798622f..cabec2839b2 100644 --- a/ax/service/utils/instantiation.py +++ b/ax/service/utils/instantiation.py @@ -501,7 +501,7 @@ def outcome_constraint_from_str( f"Outcome constraint bound should be a float for '{representation}'." ) return OutcomeConstraint( - cls._make_metric( + metric=cls._make_metric( name=tokens[0], for_opt_config=True, metric_definitions=metric_definitions, @@ -521,12 +521,11 @@ def objective_threshold_constraint_from_str( oc = cls.outcome_constraint_from_str( representation, metric_definitions=metric_definitions ) - return ObjectiveThreshold( - metric=oc.metric.clone(), - bound=oc.bound, - relative=oc.relative, - op=oc.op, - ) + # Create an ObjectiveThreshold that shares the same expression string + # as the OutcomeConstraint, bypassing the deprecated __init__. + ot = ObjectiveThreshold.__new__(ObjectiveThreshold) + ot._expression_str = oc._expression_str + return ot @classmethod def make_objectives( @@ -586,8 +585,10 @@ def make_objective_thresholds( objective_thresholds: list[str], status_quo_defined: bool, metric_definitions: dict[str, dict[str, Any]] | None = None, - ) -> list[ObjectiveThreshold]: - typed_objective_thresholds = ( + ) -> list[OutcomeConstraint]: + # pyre-ignore[9]: ObjectiveThreshold is a subclass of OutcomeConstraint; + # list invariance prevents direct assignment. + typed_objective_thresholds: list[OutcomeConstraint] = ( [ cls.objective_threshold_constraint_from_str( c, metric_definitions=metric_definitions @@ -610,7 +611,7 @@ def make_objective_thresholds( @staticmethod def optimization_config_from_objectives( objectives: list[Objective], - objective_thresholds: list[ObjectiveThreshold], + objective_thresholds: list[OutcomeConstraint], outcome_constraints: list[OutcomeConstraint], ) -> OptimizationConfig: """Parse objectives and constraints to define optimization config. @@ -882,6 +883,32 @@ def make_experiment( ] ) + # Also create properly-typed metrics for optimization config metric + # names so they are registered with proper types and properties + # before the auto-registration in Experiment.__init__ fires. + if optimization_config is not None: + # Build lower_is_better map from the optimization config so + # that pre-registered metrics carry the correct directionality. + lower_is_better_map: dict[str, bool] = {} + for mn, weight in optimization_config.objective.metric_weights: + lower_is_better_map[mn] = weight < 0 + for constraint in optimization_config.outcome_constraints: + for mn, _ in constraint.metric_weights: + if mn not in lower_is_better_map: + lower_is_better_map[mn] = constraint.op is ComparisonOp.LEQ + + tracking_names = {m.name for m in (tracking_metrics or [])} + for metric_name in optimization_config.metric_names: + if metric_name not in tracking_names: + opt_metric = cls._make_metric( + name=metric_name, + metric_definitions=metric_definitions, + lower_is_better=lower_is_better_map.get(metric_name), + ) + if tracking_metrics is None: + tracking_metrics = [] + tracking_metrics.append(opt_metric) + properties: dict[str, Any] = {} if immutable_search_space_and_opt_config: diff --git a/ax/service/utils/report_utils.py b/ax/service/utils/report_utils.py index 274507ea464..53431084285 100644 --- a/ax/service/utils/report_utils.py +++ b/ax/service/utils/report_utils.py @@ -33,7 +33,7 @@ from ax.core.map_metric import MapMetric from ax.core.metric import Metric from ax.core.multi_type_experiment import MultiTypeExperiment -from ax.core.objective import MultiObjective, ScalarizedObjective +from ax.core.objective import ScalarizedObjective from ax.core.optimization_config import ( MultiObjectiveOptimizationConfig, OptimizationConfig, @@ -126,7 +126,7 @@ def _get_objective_trace_plot( metric_names = ( metric_name for metric_name in [ - optimization_config.objective.metric.name, + optimization_config.objective.metric_names[0], true_objective_metric_name, ] if metric_name is not None @@ -138,7 +138,7 @@ def _get_objective_trace_plot( metric_colname=metric_name, minimize=none_throws( optimization_config.objective.minimize - if optimization_config.objective.metric.name == metric_name + if optimization_config.objective.metric_names[0] == metric_name else experiment.metrics[metric_name].lower_is_better ), title=f"Best {metric_name} found vs. trial index", @@ -441,7 +441,7 @@ def get_standard_plots( if map_metrics: # Sort so that objective metrics appear first map_metrics.sort( - key=lambda e: e.name in [m.name for m in objective.metrics], + key=lambda e: e.name in objective.metric_names, reverse=True, ) for by_walltime in [False, True]: @@ -994,10 +994,8 @@ def _get_metric_name_pairs( experiment=experiment ) if none_throws(optimization_config).is_moo_problem: - multi_objective = assert_is_instance( - none_throws(optimization_config).objective, MultiObjective - ) - metric_names = [obj.metric.name for obj in multi_objective.objectives] + objective = none_throws(optimization_config).objective + metric_names = list(objective.metric_names) if len(metric_names) > use_first_n_metrics: logger.info( f"Got `metric_names = {metric_names}` of length {len(metric_names)}. " @@ -1008,9 +1006,9 @@ def _get_metric_name_pairs( metric_name_pairs = itertools.combinations(metric_names, 2) return metric_name_pairs raise UserInputError( - "Inference of `metric_names` failed. Expected `MultiObjective` but " + "Inference of `metric_names` failed. Expected a multi-objective but " f"got {none_throws(optimization_config).objective}. Please provide an " - "experiment with a MultiObjective `optimization_config`." + "experiment with a multi-objective `optimization_config`." ) @@ -1340,32 +1338,31 @@ def maybe_extract_baseline_comparison_values( baseline_rows = arms_df[arms_df["arm_name"] == baseline_arm_name] if experiment.is_moo_problem: - multi_objective = assert_is_instance( - optimization_config.objective, MultiObjective - ) + objective = optimization_config.objective result_list = [] - for objective in multi_objective.objectives: - name = objective.metric.name - minimize = objective.minimize + for metric_name, weight in objective.metric_weights: + minimize = weight < 0 # Check if metric column exists in both comparison and baseline dataframes if ( - name not in comparison_arm_df.columns - or name not in baseline_rows.columns + metric_name not in comparison_arm_df.columns + or metric_name not in baseline_rows.columns ): - logger.debug(f"compare_to_baseline: metric '{name}' not found in data.") + logger.debug( + f"compare_to_baseline: metric '{metric_name}' not found in data." + ) return None opt_index = ( - comparison_arm_df[name].idxmin() + comparison_arm_df[metric_name].idxmin() if minimize - else comparison_arm_df[name].idxmax() + else comparison_arm_df[metric_name].idxmax() ) comparison_row = arms_df.iloc[opt_index] - baseline_value = baseline_rows.iloc[0][name] + baseline_value = baseline_rows.iloc[0][metric_name] result_tuple = _build_result_tuple( - objective_name=name, + objective_name=metric_name, objective_minimize=minimize, baseline_arm_name=baseline_arm_name, baseline_value=baseline_value, @@ -1375,7 +1372,7 @@ def maybe_extract_baseline_comparison_values( result_list.append(result_tuple) return result_list if result_list else None - objective_name = optimization_config.objective.metric.name + objective_name = optimization_config.objective.metric_names[0] # Check if metric column exists in both comparison and baseline dataframes if ( @@ -1473,7 +1470,7 @@ def warn_if_unpredictable_metrics( metric_names = list(experiment.metrics.keys()) else: metric_names = list( - none_throws(experiment.optimization_config).metrics.keys() + none_throws(experiment.optimization_config).metric_names ) else: # Raise a ValueError if any metric names are invalid. @@ -1497,10 +1494,7 @@ def warn_if_unpredictable_metrics( def _has_reference_point(optimization_config: MultiObjectiveOptimizationConfig) -> bool: - objectives = assert_is_instance( - optimization_config.objective, MultiObjective - ).objectives - objective_names = {obj.metric.name for obj in objectives} + objective_names = set(optimization_config.objective.metric_names) thresholds = optimization_config.objective_thresholds - threshold_names = {threshold.metric.name for threshold in thresholds} + threshold_names = {threshold.metric_names[0] for threshold in thresholds} return objective_names == threshold_names diff --git a/ax/storage/json_store/decoder.py b/ax/storage/json_store/decoder.py index 3cc658b51e3..fc5827a89a1 100644 --- a/ax/storage/json_store/decoder.py +++ b/ax/storage/json_store/decoder.py @@ -715,7 +715,7 @@ def multi_type_experiment_from_json( experiment = MultiTypeExperiment(**kwargs) for metric in tracking_metrics: - experiment._tracking_metrics[metric.name] = metric + experiment._metrics[metric.name] = metric experiment._metric_to_canonical_name = _metric_to_canonical_name experiment._metric_to_trial_type = _metric_to_trial_type experiment._trial_type_to_runner = _trial_type_to_runner @@ -1445,10 +1445,13 @@ def objective_from_json( ) -> Objective: """Load an ``Objective`` from JSON in a backwards compatible way. - If both ``minimize`` and ``lower_is_better`` are specified but have conflicting - values, this will overwrite ``lower_is_better=minimize`` to resolve the conflict. + Supports two formats: + - New expression-based format: ``{"expression": "..."}`` + - Legacy format: ``{"metric": {...}, "minimize": true/false}`` - # TODO: Do we need to do this for scalarized objective as well? + For the legacy format, if both ``minimize`` and ``lower_is_better`` are + specified but have conflicting values, this will overwrite + ``lower_is_better=minimize`` to resolve the conflict. """ input_args = { k: object_from_json( @@ -1458,6 +1461,10 @@ def objective_from_json( ) for k, v in object_json.items() } + # New expression-based format + if "expression" in input_args: + return Objective(expression=input_args["expression"]) + # Legacy metric/minimize format metric = input_args.pop("metric") minimize = input_args.pop("minimize") if metric.lower_is_better is not None and metric.lower_is_better != minimize: diff --git a/ax/storage/json_store/encoders.py b/ax/storage/json_store/encoders.py index fe62efec79b..39eefc90db8 100644 --- a/ax/storage/json_store/encoders.py +++ b/ax/storage/json_store/encoders.py @@ -106,14 +106,21 @@ def analysis_card_group_to_dict(group: AnalysisCardGroup) -> dict[str, Any]: def experiment_to_dict(experiment: Experiment) -> dict[str, Any]: """Convert Ax experiment to a dictionary.""" + opt_config = experiment.optimization_config return { "__type": experiment.__class__.__name__, "name": experiment._name, "description": experiment.description, "experiment_type": experiment.experiment_type, "search_space": experiment.search_space, - "optimization_config": experiment.optimization_config, - "tracking_metrics": list(experiment._tracking_metrics.values()), + "optimization_config": ( + _build_opt_config_dict( + opt_config=opt_config, experiment_metrics=experiment._metrics + ) + if opt_config is not None + else None + ), + "tracking_metrics": list(experiment.metrics.values()), "runner": experiment.runner, "status_quo": experiment.status_quo, "time_created": experiment.time_created, @@ -290,36 +297,65 @@ def metric_to_dict(metric: Metric) -> dict[str, Any]: def objective_to_dict(objective: Objective) -> dict[str, Any]: """Convert Ax objective to a dictionary.""" + # Plain Objective: serialize using expression string. + if type(objective) is Objective: + return {"__type": "Objective", "expression": objective.expression} + metric_name = objective.metric_names[0] + minimize = objective.minimize return { "__type": objective.__class__.__name__, - "metric": objective.metric, - "minimize": objective.minimize, + "metric": Metric(name=metric_name, lower_is_better=minimize), + "minimize": minimize, } def multi_objective_to_dict(objective: MultiObjective) -> dict[str, Any]: - """Convert Ax objective to a dictionary.""" + """Convert Ax multi-objective to a dictionary.""" + sub_objectives = [] + for name, weight in objective.metric_weights: + minimize = weight < 0 + sub_objectives.append( + Objective( + metric=Metric(name=name, lower_is_better=minimize), + minimize=minimize, + ) + ) return { "__type": objective.__class__.__name__, - "objectives": objective.objectives, + "objectives": sub_objectives, } def scalarized_objective_to_dict(objective: ScalarizedObjective) -> dict[str, Any]: - """Convert Ax objective to a dictionary.""" + """Convert Ax scalarized objective to a dictionary.""" + metric_weights = objective.metric_weights + # Infer minimize from effective weights: if all weights are <= 0 (with at + # least one < 0), the objective was constructed with minimize=True, which + # negates all raw weights. + minimize = all(w <= 0 for _, w in metric_weights) and any( + w < 0 for _, w in metric_weights + ) + metrics = [] + weights = [] + for name, w in metric_weights: + metrics.append(Metric(name=name)) + # metric_weights returns effective weights (sign-adjusted for minimize). + # JSON format expects raw weights + separate minimize flag. + weights.append(w if not minimize else -w) return { "__type": objective.__class__.__name__, - "metrics": objective.metrics, - "weights": objective.weights, - "minimize": objective.minimize, + "metrics": metrics, + "weights": weights, + "minimize": minimize, } def outcome_constraint_to_dict(outcome_constraint: OutcomeConstraint) -> dict[str, Any]: """Convert Ax outcome constraint to a dictionary.""" + metric_name = outcome_constraint.metric_names[0] return { "__type": outcome_constraint.__class__.__name__, - "metric": outcome_constraint.metric, + "metric": Metric(name=metric_name), "op": outcome_constraint.op, "bound": outcome_constraint.bound, "relative": outcome_constraint.relative, @@ -713,3 +749,116 @@ def backend_simulator_to_dict( "completed": backend_simulator._completed, "verbose_logging": backend_simulator._verbose_logging, } + + +# --------------------------------------------------------------------------- +# Private helpers for experiment_to_dict +# --------------------------------------------------------------------------- + + +def _build_opt_config_dict( + opt_config: OptimizationConfig, + experiment_metrics: dict[str, Metric], +) -> dict[str, Any]: + """Build optimization config dict with real metric objects from experiment. + + This bypasses the encoder registry dispatch so that real Metric objects + (e.g. ``BraninMetric``) are preserved as leaf values. ``object_to_json`` + will then recursively encode them via ``metric_to_dict``, capturing the + full metric type. + """ + objective_dict = _build_objective_dict( + objective=opt_config.objective, experiment_metrics=experiment_metrics + ) + constraint_dicts = [ + _build_constraint_dict(constraint=c, experiment_metrics=experiment_metrics) + for c in opt_config.outcome_constraints + ] + result: dict[str, Any] = { + "__type": opt_config.__class__.__name__, + "objective": objective_dict, + "outcome_constraints": constraint_dicts, + "pruning_target_parameterization": opt_config.pruning_target_parameterization, + } + if isinstance(opt_config, MultiObjectiveOptimizationConfig): + result["objective_thresholds"] = [ + _build_constraint_dict(constraint=t, experiment_metrics=experiment_metrics) + for t in opt_config.objective_thresholds + ] + if isinstance(opt_config, PreferenceOptimizationConfig): + result[PREFERENCE_PROFILE_NAME] = opt_config.preference_profile_name + result[EXPECT_RELATIVIZED_OUTCOMES] = opt_config.expect_relativized_outcomes + return result + + +def _build_objective_dict( + objective: Objective, + experiment_metrics: dict[str, Metric], +) -> dict[str, Any]: + """Build objective dict with real metric objects from experiment.""" + # Plain Objective (not a deprecated subclass like ScalarizedObjective or + # MultiObjective): serialize using the expression string directly. + if type(objective) is Objective: + return {"__type": "Objective", "expression": objective.expression} + if objective.is_multi_objective: + sub_objectives = [] + for name, weight in objective.metric_weights: + metric = experiment_metrics.get(name, Metric(name=name)) + minimize = weight < 0 + sub_objectives.append( + { + "__type": "Objective", + "metric": metric, + "minimize": minimize, + } + ) + return { + "__type": objective.__class__.__name__, + "objectives": sub_objectives, + } + elif objective.is_scalarized_objective: + metric_weights = objective.metric_weights + # Infer minimize from effective weights: if all weights are <= 0 (with + # at least one < 0), the objective was constructed with minimize=True. + minimize = all(w <= 0 for _, w in metric_weights) and any( + w < 0 for _, w in metric_weights + ) + metrics = [] + weights = [] + for name, w in metric_weights: + metric = experiment_metrics.get(name, Metric(name=name)) + metrics.append(metric) + # metric_weights returns effective weights (sign-adjusted for + # minimize). JSON format expects raw weights + separate minimize + # flag, so undo the sign adjustment. + weights.append(w if not minimize else -w) + return { + "__type": objective.__class__.__name__, + "metrics": metrics, + "weights": weights, + "minimize": minimize, + } + else: + metric_name = objective.metric_names[0] + metric = experiment_metrics.get(metric_name, Metric(name=metric_name)) + return { + "__type": objective.__class__.__name__, + "metric": metric, + "minimize": objective.minimize, + } + + +def _build_constraint_dict( + constraint: OutcomeConstraint, + experiment_metrics: dict[str, Metric], +) -> dict[str, Any]: + """Build constraint dict with real metric objects from experiment.""" + metric_name = constraint.metric_names[0] + metric = experiment_metrics.get(metric_name, Metric(name=metric_name)) + return { + "__type": constraint.__class__.__name__, + "metric": metric, + "op": constraint.op, + "bound": constraint.bound, + "relative": constraint.relative, + } diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py index 38993cba4b4..0c8efabb649 100644 --- a/ax/storage/json_store/tests/test_json_store.py +++ b/ax/storage/json_store/tests/test_json_store.py @@ -1141,16 +1141,20 @@ def test_parameter_constraint_backwards_compatibility(self) -> None: def test_objective_backwards_compatibility(self) -> None: # Test that we can load an objective that has conflicting # ``lower_is_better`` and ``minimize`` fields. - objective = get_objective(minimize=True) - objective.metric.lower_is_better = False # for conflict! - objective_json = object_to_json(objective) - self.assertTrue(objective_json["minimize"]) - self.assertFalse(objective_json["metric"]["lower_is_better"]) + # In the old format, the metric had ``lower_is_better`` and the + # objective had ``minimize``. When they disagree, ``minimize`` wins. + objective_json = { + "__type": "Objective", + "metric": { + "__type": "Metric", + "name": "m1", + "lower_is_better": False, + }, + "minimize": True, + } objective_loaded = object_from_json(objective_json) self.assertIsInstance(objective_loaded, Objective) - self.assertNotEqual(objective, objective_loaded) self.assertTrue(objective_loaded.minimize) - self.assertTrue(objective_loaded.metric.lower_is_better) def test_generation_step_backwards_compatibility(self) -> None: # Test that we can load a generation step with deprecated kwargs. diff --git a/ax/storage/sqa_store/decoder.py b/ax/storage/sqa_store/decoder.py index ec138e36ee0..3a9c043f1d8 100644 --- a/ax/storage/sqa_store/decoder.py +++ b/ax/storage/sqa_store/decoder.py @@ -8,6 +8,7 @@ from collections import defaultdict from collections.abc import Callable +from copy import deepcopy from enum import Enum from io import StringIO from logging import Logger @@ -220,13 +221,16 @@ def _init_experiment_from_sqa( # `experiment_sqa.properties` is `sqlalchemy.ext.mutable.MutableDict` # so need to convert it to regular dict. properties = dict(experiment_sqa.properties or {}) - opt_config, tracking_metrics = self.opt_config_and_tracking_metrics_from_sqa( - metrics_sqa=experiment_sqa.metrics, - pruning_target_parameterization=( - self._get_pruning_target_parameterization_from_experiment_properties( - properties=properties - ) - ), + pruning_target = ( + self._get_pruning_target_parameterization_from_experiment_properties( + properties=properties + ) + ) + opt_config, _tracking_metrics, all_metrics = ( + self.opt_config_and_tracking_metrics_from_sqa( + metrics_sqa=experiment_sqa.metrics, + pruning_target_parameterization=pruning_target, + ) ) search_space = self.search_space_from_sqa( parameters_sqa=experiment_sqa.parameters, @@ -267,7 +271,7 @@ def _init_experiment_from_sqa( description=experiment_sqa.description, search_space=search_space, optimization_config=opt_config, - tracking_metrics=tracking_metrics, + tracking_metrics=all_metrics, runner=runner, status_quo=status_quo, is_test=experiment_sqa.is_test, @@ -281,13 +285,16 @@ def _init_mt_experiment_from_sqa( ) -> MultiTypeExperiment: """First step of conversion within experiment_from_sqa.""" properties = dict(experiment_sqa.properties or {}) - opt_config, tracking_metrics = self.opt_config_and_tracking_metrics_from_sqa( - metrics_sqa=experiment_sqa.metrics, - pruning_target_parameterization=( - self._get_pruning_target_parameterization_from_experiment_properties( - properties=properties - ) - ), + pruning_target = ( + self._get_pruning_target_parameterization_from_experiment_properties( + properties=properties + ) + ) + opt_config, tracking_metrics, all_metrics = ( + self.opt_config_and_tracking_metrics_from_sqa( + metrics_sqa=experiment_sqa.metrics, + pruning_target_parameterization=pruning_target, + ) ) search_space = self.search_space_from_sqa( parameters_sqa=experiment_sqa.parameters, @@ -331,6 +338,16 @@ def _init_mt_experiment_from_sqa( status_quo=status_quo, properties=properties, ) + # Replace any placeholder Metric objects (created by __init__'s + # auto-registration for optimization config metric names) with the + # properly typed metrics decoded from the database (e.g. BraninMetric). + for metric in all_metrics: + if ( + metric.name in experiment._metrics + and type(experiment._metrics[metric.name]) is Metric + and type(metric) is not Metric + ): + experiment._metrics[metric.name] = metric # pyre-fixme[8]: `_trial_type_to_runner` expects `Dict[Optional[str], # Optional[Runner]]` but the dict built here uses `str` keys. experiment._trial_type_to_runner = trial_type_to_runner @@ -601,32 +618,109 @@ def metric_from_sqa( def opt_config_and_tracking_metrics_from_sqa( self, metrics_sqa: list[SQAMetric], pruning_target_parameterization: Arm | None - ) -> tuple[OptimizationConfig | None, list[Metric]]: + ) -> tuple[OptimizationConfig | None, list[Metric], list[Metric]]: """Convert a list of SQLAlchemy Metrics to Ax OptimizationConfig and tracking metrics. + + Returns: + A tuple of (optimization_config, tracking_metrics, all_metrics). + ``tracking_metrics`` contains only non-optimization metrics. + ``all_metrics`` contains all decoded Metric objects, including + those used in the optimization config, so the Experiment can + register the full metric types (e.g. BraninMetric) rather than + plain Metric placeholders. """ objective = None objective_thresholds = [] outcome_constraints = [] tracking_metrics = [] preference_objective_sqa = None + # Collect all decoded Metric objects (including those used in the + # optimization config) so the Experiment can register the real metric + # types (e.g. BraninMetric) rather than plain Metric placeholders. + all_metrics: list[Metric] = [] for metric_sqa in metrics_sqa: if metric_sqa.intent == MetricIntent.PREFERENCE_OBJECTIVE: preference_objective_sqa = metric_sqa - metric = self.metric_from_sqa(metric_sqa=metric_sqa) - if isinstance(metric, Objective): - objective = metric - elif isinstance(metric, ObjectiveThreshold): - objective_thresholds.append(metric) - elif isinstance(metric, OutcomeConstraint): - outcome_constraints.append(metric) + # Decode the raw Metric first (before wrapping in Objective etc.) + raw_metric = self._metric_from_sqa_util(metric_sqa) + + result = self.metric_from_sqa(metric_sqa=metric_sqa) + if isinstance(result, Objective): + objective = result + # Collect metrics from the objective + if metric_sqa.intent in ( + MetricIntent.MULTI_OBJECTIVE, + MetricIntent.PREFERENCE_OBJECTIVE, + MetricIntent.SCALARIZED_OBJECTIVE, + ): + # For multi/scalarized objectives, decode each child metric + try: + children_sqa = ( + metric_sqa.scalarized_objective_children_metrics or [] + ) + except DetachedInstanceError: + children_sqa = _get_scalarized_objective_children_metrics( + metric_sqa.id, self + ) + # Apply skip_runners_and_metrics to children if set + if metric_sqa.properties and metric_sqa.properties.get( + "skip_runners_and_metrics" + ): + for child_sqa in children_sqa: + child_sqa.metric_type = self.config.metric_registry[Metric] + for child_sqa in children_sqa: + child_metric = self._metric_from_sqa_util(child_sqa) + # Clear db_id: child SQA rows have + # scalarized_objective_id set (not experiment_id), + # so their IDs must not leak into experiment._metrics + # which are matched against experiment-level SQA rows. + child_metric._db_id = None + all_metrics.append(child_metric) + else: + all_metrics.append(raw_metric) + elif isinstance(result, ObjectiveThreshold): + objective_thresholds.append(result) + all_metrics.append(raw_metric) + elif isinstance(result, OutcomeConstraint): + outcome_constraints.append(result) + if metric_sqa.intent == MetricIntent.SCALARIZED_OUTCOME_CONSTRAINT: + # For scalarized outcome constraints, decode each child + # metric rather than the parent placeholder. + try: + children_sqa = ( + metric_sqa.scalarized_outcome_constraint_children_metrics + or [] + ) + except DetachedInstanceError: + children_sqa = ( + _get_scalarized_outcome_constraint_children_metrics( + metric_sqa.id, self + ) + ) + if metric_sqa.properties and metric_sqa.properties.get( + "skip_runners_and_metrics" + ): + for child_sqa in children_sqa: + child_sqa.metric_type = self.config.metric_registry[Metric] + for child_sqa in children_sqa: + child_metric = self._metric_from_sqa_util(child_sqa) + # Clear db_id: child SQA rows have + # scalarized_outcome_constraint_id set (not + # experiment_id), so their IDs must not leak into + # experiment._metrics. + child_metric._db_id = None + all_metrics.append(child_metric) + else: + all_metrics.append(raw_metric) else: - tracking_metrics.append(metric) + tracking_metrics.append(result) + all_metrics.append(raw_metric) if objective is None: - return None, tracking_metrics + return None, tracking_metrics, all_metrics if preference_objective_sqa is not None: if objective_thresholds: @@ -658,7 +752,7 @@ def opt_config_and_tracking_metrics_from_sqa( outcome_constraints=outcome_constraints, pruning_target_parameterization=pruning_target_parameterization, ) - return (optimization_config, tracking_metrics) + return (optimization_config, tracking_metrics, all_metrics) def arm_from_sqa(self, arm_sqa: SQAArm) -> Arm: """Convert SQLAlchemy Arm to Ax Arm.""" @@ -713,6 +807,7 @@ def generator_run_from_sqa( ( opt_config, tracking_metrics, + _all_metrics, ) = self.opt_config_and_tracking_metrics_from_sqa( metrics_sqa=generator_run_sqa.metrics, pruning_target_parameterization=None, @@ -1137,7 +1232,7 @@ def _metric_from_sqa_util(self, metric_sqa: SQAMetric) -> Metric: args = dict( object_from_json( - metric_sqa.properties, + deepcopy(metric_sqa.properties or {}), decoder_registry=self.config.json_decoder_registry, class_decoder_registry=self.config.json_class_decoder_registry, ) @@ -1205,6 +1300,14 @@ def _scalarized_objective_from_sqa(self, parent_metric_sqa: SQAMetric) -> Object "because minimize is None." ) + # If the parent metric has an expression in its properties, use it + # directly to create an Objective (preserving the original expression). + parent_properties = parent_metric_sqa.properties or {} + if "expression" in parent_properties: + obj = Objective(expression=parent_properties["expression"]) + obj.db_id = parent_metric_sqa.id + return obj + metrics_sqa_children = self._get_and_process_children_metrics( parent_metric_sqa=parent_metric_sqa, children_attribute_name="scalarized_objective_children_metrics", @@ -1242,12 +1345,14 @@ def _outcome_constraint_from_sqa( "Cannot decode SQAMetric to OutcomeConstraint because " "bound, op, or relative is None." ) - return OutcomeConstraint( + oc = OutcomeConstraint( metric=metric, bound=float(none_throws(metric_sqa.bound)), op=none_throws(metric_sqa.op), relative=none_throws(metric_sqa.relative), ) + oc.db_id = metric_sqa.id + return oc def _scalarized_outcome_constraint_from_sqa( self, metric: Metric, metric_sqa: SQAMetric @@ -1311,9 +1416,7 @@ def _objective_threshold_from_sqa( relative=relative, op=metric_sqa.op, ) - # ObjectiveThreshold constructor clones the passed-in metric, which means - # the db id gets lost and so we need to reset it - ot.metric._db_id = metric.db_id + ot.db_id = metric_sqa.id return ot def _get_and_process_children_metrics( diff --git a/ax/storage/sqa_store/encoder.py b/ax/storage/sqa_store/encoder.py index 6c99add37b4..f5b552c973f 100644 --- a/ax/storage/sqa_store/encoder.py +++ b/ax/storage/sqa_store/encoder.py @@ -165,11 +165,12 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment: """ optimization_metrics = self.optimization_config_to_sqa( - experiment.optimization_config + experiment.optimization_config, + experiment_metrics=experiment._metrics, ) tracking_metrics = [] - for metric in experiment._tracking_metrics.values(): + for metric in experiment.tracking_metrics: tracking_metrics.append(self.metric_to_sqa(metric)) parameters, parameter_constraints = self.search_space_to_sqa( @@ -184,7 +185,10 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment: trials = [] for trial in experiment.trials.values(): - trial_sqa = self.trial_to_sqa(trial=trial) + trial_sqa = self.trial_to_sqa( + trial=trial, + experiment_metrics=experiment._metrics, + ) trials.append(trial_sqa) experiment_data = self.experiment_data_to_sqa(experiment=experiment) @@ -438,35 +442,53 @@ def get_children_metrics_by_name( for (metric, weight) in zip(metrics, weights) } - def objective_to_sqa(self, objective: Objective) -> SQAMetric: + def objective_to_sqa( + self, + objective: Objective, + experiment_metrics: dict[str, Metric] | None = None, + ) -> SQAMetric: """Convert Ax Objective to SQLAlchemy.""" - if isinstance(objective, ScalarizedObjective): - objective_sqa = self.scalarized_objective_to_sqa(objective) + if objective.is_scalarized_objective or isinstance( + objective, ScalarizedObjective + ): + objective_sqa = self.scalarized_objective_to_sqa( + objective, experiment_metrics=experiment_metrics + ) - elif isinstance(objective, MultiObjective): - objective_sqa = self.multi_objective_to_sqa(objective) + elif objective.is_multi_objective or isinstance(objective, MultiObjective): + objective_sqa = self.multi_objective_to_sqa( + objective, experiment_metrics=experiment_metrics + ) else: - metric_type, properties = self.get_metric_type_and_properties( - metric=objective.metric + metric_name = objective.metric_names[0] + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) ) + metric_type, properties = self.get_metric_type_and_properties(metric=metric) metric_class = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) objective_sqa = ( metric_class( # pyre-ignore[29]: `SQAMetric` is not a function. - id=objective.metric.db_id, - name=objective.metric.name, - signature=objective.metric.signature, + id=objective.db_id, + name=metric.name, + signature=metric.signature, metric_type=metric_type, intent=MetricIntent.OBJECTIVE, minimize=objective.minimize, properties=properties, - lower_is_better=objective.metric.lower_is_better, + lower_is_better=metric.lower_is_better, ) ) return assert_is_instance(objective_sqa, SQAMetric) - def multi_objective_to_sqa(self, multi_objective: MultiObjective) -> SQAMetric: + def multi_objective_to_sqa( + self, + multi_objective: Objective, + experiment_metrics: dict[str, Metric] | None = None, + ) -> SQAMetric: """Convert Ax Multi Objective to SQLAlchemy. Returns: A parent `SQAMetric`, whose children are the `SQAMetric`-s @@ -476,21 +498,26 @@ def multi_objective_to_sqa(self, multi_objective: MultiObjective) -> SQAMetric: # Constructing children SQAMetric classes (these are the real metrics in # the `MultiObjective`). children_objectives = [] - for objective in multi_objective.objectives: - objective_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) - type_and_properties = self.get_metric_type_and_properties( - metric=objective.metric + metric_weights_dict = dict(multi_objective.metric_weights) + for metric_name in multi_objective.metric_names: + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) ) + objective_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) + type_and_properties = self.get_metric_type_and_properties(metric=metric) + weight = metric_weights_dict.get(metric_name, 1.0) children_objectives.append( objective_cls( # pyre-ignore[29]: `SQAMetric` is not a func. - id=objective.metric.db_id, - name=objective.metric.name, - signature=objective.metric.signature, + id=None, + name=metric.name, + signature=metric.signature, metric_type=type_and_properties[0], intent=MetricIntent.OBJECTIVE, - minimize=objective.minimize, + minimize=weight < 0, properties=type_and_properties[1], - lower_is_better=objective.metric.lower_is_better, + lower_is_better=metric.lower_is_better, ) ) @@ -511,9 +538,10 @@ def multi_objective_to_sqa(self, multi_objective: MultiObjective) -> SQAMetric: def preference_objective_to_sqa( self, - multi_objective: MultiObjective, + multi_objective: Objective, preference_profile_name: str, expect_relativized_outcomes: bool, + experiment_metrics: dict[str, Metric] | None = None, ) -> SQAMetric: """Convert Ax PreferenceOptimizationConfig objective to SQLAlchemy. @@ -527,21 +555,26 @@ def preference_objective_to_sqa( # Constructing children SQAMetric classes (these are the real metrics in # the `MultiObjective`). children_objectives = [] - for objective in multi_objective.objectives: - objective_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) - type_and_properties = self.get_metric_type_and_properties( - metric=objective.metric + metric_weights_dict = dict(multi_objective.metric_weights) + for metric_name in multi_objective.metric_names: + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) ) + objective_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) + type_and_properties = self.get_metric_type_and_properties(metric=metric) + weight = metric_weights_dict.get(metric_name, 1.0) children_objectives.append( objective_cls( # pyre-ignore[29]: `SQAMetric` is not a func. - id=objective.metric.db_id, - name=objective.metric.name, - signature=objective.metric.signature, + id=None, + name=metric.name, + signature=metric.signature, metric_type=type_and_properties[0], intent=MetricIntent.OBJECTIVE, - minimize=objective.minimize, + minimize=weight < 0, properties=type_and_properties[1], - lower_is_better=objective.metric.lower_is_better, + lower_is_better=metric.lower_is_better, ) ) @@ -567,42 +600,62 @@ def preference_objective_to_sqa( ) return parent_metric - def scalarized_objective_to_sqa(self, objective: ScalarizedObjective) -> SQAMetric: + def scalarized_objective_to_sqa( + self, + objective: Objective, + experiment_metrics: dict[str, Metric] | None = None, + ) -> SQAMetric: """Convert Ax Scalarized Objective to SQLAlchemy. Returns: A parent `SQAMetric`, whose children are the `SQAMetric`-s corresponding to `metrics` attribute of `ScalarizedObjective`. NOTE: The parent is used as a placeholder for storage purposes. """ - metrics, weights = objective.metrics, objective.weights - if (not (metrics and weights)) or len(metrics) != len(weights): + metric_weights = objective.metric_weights + if not metric_weights: raise SQAEncodeError( "Metrics and weights in scalarized objective " "must be lists of equal length." ) - metrics_by_name = self.get_children_metrics_by_name( - metrics=metrics, weights=weights + + # Infer minimize from effective weights: if all weights are <= 0 (with + # at least one < 0), the objective was constructed with minimize=True. + minimize = all(w <= 0 for _, w in metric_weights) and any( + w < 0 for _, w in metric_weights ) # Constructing children SQAMetric classes (these are the real metrics in # the `ScalarizedObjective`). children_metrics = [] - for metric_name in metrics_by_name: - m, w, metric_cls, type_and_properties = metrics_by_name[metric_name] + for metric_name, w in metric_weights: + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) + ) + metric_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) + type_and_properties = self.get_metric_type_and_properties(metric=metric) children_metrics.append( metric_cls( # pyre-ignore[29]: `SQAMetric` is not a function. - id=m.db_id, + id=None, name=metric_name, - signature=m.signature, + signature=metric.signature, metric_type=type_and_properties[0], intent=MetricIntent.OBJECTIVE, - minimize=objective.minimize, + minimize=minimize, properties=type_and_properties[1], - lower_is_better=m.lower_is_better, + lower_is_better=metric.lower_is_better, scalarized_objective_weight=w, ) ) + # For plain Objective instances (not deprecated ScalarizedObjective + # subclass), store the original expression so it can be restored + # losslessly during decoding. + parent_properties: dict[str, Any] = {} + if type(objective) is Objective: + parent_properties["expression"] = objective.expression + # Constructing a parent SQAMetric class parent_metric_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) parent_metric = parent_metric_cls( # pyre-ignore[29]: `SQAMetric` not a func. @@ -610,28 +663,38 @@ def scalarized_objective_to_sqa(self, objective: ScalarizedObjective) -> SQAMetr name="scalarized_objective", metric_type=self.config.metric_registry[Metric], intent=MetricIntent.SCALARIZED_OBJECTIVE, - minimize=objective.minimize, - lower_is_better=objective.minimize, + minimize=minimize, + lower_is_better=minimize, scalarized_objective_children_metrics=children_metrics, signature="scalarized_objective", + properties=parent_properties, ) return parent_metric def outcome_constraint_to_sqa( - self, outcome_constraint: OutcomeConstraint + self, + outcome_constraint: OutcomeConstraint, + experiment_metrics: dict[str, Metric] | None = None, ) -> SQAMetric: """Convert Ax OutcomeConstraint to SQLAlchemy.""" if isinstance(outcome_constraint, ScalarizedOutcomeConstraint): - return self.scalarized_outcome_constraint_to_sqa(outcome_constraint) + return self.scalarized_outcome_constraint_to_sqa( + outcome_constraint, experiment_metrics=experiment_metrics + ) - metric = outcome_constraint.metric + metric_name = outcome_constraint.metric_names[0] + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) + ) metric_type, properties = self.get_metric_type_and_properties(metric=metric) # pyre-fixme: Expected `Base` for 1st...t `typing.Type[Metric]`. metric_class: SQAMetric = self.config.class_to_sqa_class[Metric] # pyre-fixme[29]: `SQAMetric` is not a function. constraint_sqa = metric_class( - id=metric.db_id, + id=outcome_constraint.db_id, name=metric.name, signature=metric.signature, metric_type=metric_type, @@ -645,34 +708,38 @@ def outcome_constraint_to_sqa( return constraint_sqa def scalarized_outcome_constraint_to_sqa( - self, outcome_constraint: ScalarizedOutcomeConstraint + self, + outcome_constraint: ScalarizedOutcomeConstraint, + experiment_metrics: dict[str, Metric] | None = None, ) -> SQAMetric: """Convert Ax Scalarized OutcomeConstraint to SQLAlchemy.""" - metrics, weights = outcome_constraint.metrics, outcome_constraint.weights - - if metrics is None or weights is None or len(metrics) != len(weights): + metric_weights = outcome_constraint.metric_weights + if not metric_weights: raise SQAEncodeError( "Metrics and weights in scalarized OutcomeConstraint " "must be lists of equal length." ) - metrics_by_name = self.get_children_metrics_by_name( - metrics=metrics, weights=weights - ) # Constructing children SQAMetric classes (these are the real metrics in # the `ScalarizedObjective`). children_metrics = [] - for metric_name in metrics_by_name: - m, w, metric_cls, type_and_properties = metrics_by_name[metric_name] + for metric_name, w in metric_weights: + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) + ) + metric_cls = cast(SQAMetric, self.config.class_to_sqa_class[Metric]) + type_and_properties = self.get_metric_type_and_properties(metric=metric) children_metrics.append( metric_cls( # pyre-ignore[29]: `SQAMetric` is not a function. - id=m.db_id, + id=None, name=metric_name, - signature=m.signature, + signature=metric.signature, metric_type=type_and_properties[0], intent=MetricIntent.OUTCOME_CONSTRAINT, properties=type_and_properties[1], - lower_is_better=m.lower_is_better, + lower_is_better=metric.lower_is_better, scalarized_outcome_constraint_weight=w, bound=outcome_constraint.bound, op=outcome_constraint.op, @@ -696,17 +763,24 @@ def scalarized_outcome_constraint_to_sqa( return parent_metric def objective_threshold_to_sqa( - self, objective_threshold: ObjectiveThreshold + self, + objective_threshold: ObjectiveThreshold, + experiment_metrics: dict[str, Metric] | None = None, ) -> SQAMetric: """Convert Ax OutcomeConstraint to SQLAlchemy.""" - metric = objective_threshold.metric + metric_name = objective_threshold.metric_names[0] + metric = ( + experiment_metrics.get(metric_name, Metric(name=metric_name)) + if experiment_metrics + else Metric(name=metric_name) + ) metric_type, properties = self.get_metric_type_and_properties(metric=metric) # pyre-fixme: Expected `Base` for 1st...t `typing.Type[Metric]`. metric_class: SQAMetric = self.config.class_to_sqa_class[Metric] # pyre-fixme[29]: `SQAMetric` is not a function. return metric_class( - id=metric.db_id, + id=objective_threshold.db_id, name=metric.name, signature=metric.signature, metric_type=metric_type, @@ -719,7 +793,9 @@ def objective_threshold_to_sqa( ) def optimization_config_to_sqa( - self, optimization_config: OptimizationConfig | None + self, + optimization_config: OptimizationConfig | None, + experiment_metrics: dict[str, Metric] | None = None, ) -> list[SQAMetric]: """Convert Ax OptimizationConfig to a list of SQLAlchemy Metrics.""" if optimization_config is None: @@ -730,7 +806,9 @@ def optimization_config_to_sqa( # Handle PreferenceOptimizationConfig separately if isinstance(optimization_config, PreferenceOptimizationConfig): objective = optimization_config.objective - if not isinstance(objective, MultiObjective): + if not ( + isinstance(objective, MultiObjective) or objective.is_multi_objective + ): raise SQAEncodeError( f"PreferenceOptimizationConfig requires a MultiObjective, " f"got {type(objective).__name__}" @@ -741,14 +819,19 @@ def optimization_config_to_sqa( expect_relativized_outcomes=( optimization_config.expect_relativized_outcomes ), + experiment_metrics=experiment_metrics, ) else: - obj_sqa = self.objective_to_sqa(objective=optimization_config.objective) + obj_sqa = self.objective_to_sqa( + objective=optimization_config.objective, + experiment_metrics=experiment_metrics, + ) metrics_sqa.append(obj_sqa) for constraint in optimization_config.outcome_constraints: constraint_sqa = self.outcome_constraint_to_sqa( - outcome_constraint=constraint + outcome_constraint=constraint, + experiment_metrics=experiment_metrics, ) metrics_sqa.append(constraint_sqa) if isinstance(optimization_config, MultiObjectiveOptimizationConfig) and not ( @@ -756,7 +839,10 @@ def optimization_config_to_sqa( ): for threshold in optimization_config.objective_thresholds: threshold_sqa = self.objective_threshold_to_sqa( - objective_threshold=threshold + objective_threshold=assert_is_instance( + threshold, ObjectiveThreshold + ), + experiment_metrics=experiment_metrics, ) metrics_sqa.append(threshold_sqa) return metrics_sqa @@ -789,6 +875,7 @@ def generator_run_to_sqa( generator_run: GeneratorRun, weight: float | None = None, reduced_state: bool = False, + experiment_metrics: dict[str, Metric] | None = None, ) -> SQAGeneratorRun: """Convert Ax GeneratorRun to SQLAlchemy. @@ -801,7 +888,10 @@ def generator_run_to_sqa( for arm, arm_weight in generator_run.arm_weights.items(): arms.append(self.arm_to_sqa(arm=arm, weight=arm_weight)) - metrics = self.optimization_config_to_sqa(generator_run.optimization_config) + metrics = self.optimization_config_to_sqa( + generator_run.optimization_config, + experiment_metrics=experiment_metrics, + ) parameters, parameter_constraints = self.search_space_to_sqa( generator_run.search_space ) @@ -968,7 +1058,10 @@ def runner_to_sqa(self, runner: Runner, trial_type: str | None = None) -> SQARun ) def trial_to_sqa( - self, trial: BaseTrial, generator_run_reduced_state: bool = False + self, + trial: BaseTrial, + generator_run_reduced_state: bool = False, + experiment_metrics: dict[str, Metric] | None = None, ) -> SQATrial: """Convert Ax Trial to SQLAlchemy. @@ -983,6 +1076,7 @@ def trial_to_sqa( gr_sqa = self.generator_run_to_sqa( generator_run=none_throws(trial.generator_run), reduced_state=generator_run_reduced_state, + experiment_metrics=experiment_metrics, ) generator_runs.append(gr_sqa) @@ -992,7 +1086,11 @@ def trial_to_sqa( self.abandoned_arm_to_sqa(abandoned_arm=abandoned_arm) ) for gr in trial._generator_runs: - gr_sqa = self.generator_run_to_sqa(generator_run=gr, weight=1.0) + gr_sqa = self.generator_run_to_sqa( + generator_run=gr, + weight=1.0, + experiment_metrics=experiment_metrics, + ) generator_runs.append(gr_sqa) # pyre-ignore[9]: Expected `Base` for 1st...ot `typing.Type[Trial]`. diff --git a/ax/storage/sqa_store/tests/test_sqa_store.py b/ax/storage/sqa_store/tests/test_sqa_store.py index 83878d96a0e..236de9f3c70 100644 --- a/ax/storage/sqa_store/tests/test_sqa_store.py +++ b/ax/storage/sqa_store/tests/test_sqa_store.py @@ -763,28 +763,24 @@ def test_load_experiment_skip_metrics_and_runners(self) -> None: # Check generator runs gr = trial.generator_runs[0] if composite_type == "multi_objective" and not immutable: - objectives = assert_is_instance( + multi_obj = assert_is_instance( none_throws(gr.optimization_config).objective, MultiObjective - ).objectives - for i, objective in enumerate(objectives): - metric = objective.metric - self.assertEqual(metric.name, f"m{1 + 2 * i}") - self.assertEqual(metric.signature, f"m{1 + 2 * i}") - self.assertEqual(metric.__class__, Metric) + ) + metric_names = multi_obj.metric_names + for i, name in enumerate(metric_names): + self.assertEqual(name, f"m{1 + 2 * i}") elif composite_type == "scalarized" and not immutable: # Check scalarized objective children scalarized_objective = none_throws(gr.optimization_config).objective if isinstance(scalarized_objective, ScalarizedObjective): - for metric in scalarized_objective.metrics: - self.assertEqual(metric.__class__, Metric) + self.assertTrue(len(scalarized_objective.metric_names) > 0) # Check scalarized outcome constraint children for constraint in none_throws( gr.optimization_config ).outcome_constraints: if isinstance(constraint, ScalarizedOutcomeConstraint): - for metric in constraint.metrics: - self.assertEqual(metric.__class__, Metric) + self.assertTrue(len(constraint.metric_names) > 0) @patch( f"{Decoder.__module__}.Decoder.generator_run_from_sqa", @@ -1267,8 +1263,7 @@ def test_experiment_objective_updates(self) -> None: # update objective # (should perform update in place) optimization_config = get_optimization_config() - objective = get_objective() - objective.minimize = True + objective = get_objective(minimize=True) optimization_config.objective = objective experiment.optimization_config = optimization_config save_experiment(experiment) @@ -1278,6 +1273,7 @@ def test_experiment_objective_updates(self) -> None: # replace objective # (old one should become tracking metric) + experiment.add_tracking_metric(Metric(name="objective")) optimization_config.objective = Objective( metric=Metric(name="objective"), minimize=False ) @@ -1300,8 +1296,12 @@ def test_experiment_outcome_constraint_updates(self) -> None: # update outcome constraint # (should perform update in place) optimization_config = get_optimization_config() - outcome_constraint = get_outcome_constraint() - outcome_constraint.bound = -1.0 + outcome_constraint = OutcomeConstraint( + metric=Metric(name="m2"), + op=ComparisonOp.GEQ, + bound=-1.0, + relative=True, + ) optimization_config.outcome_constraints = [outcome_constraint] experiment.optimization_config = optimization_config save_experiment(experiment) @@ -1310,6 +1310,7 @@ def test_experiment_outcome_constraint_updates(self) -> None: ) # add outcome constraint + experiment.add_tracking_metric(Metric(name="outcome")) outcome_constraint2 = OutcomeConstraint( metric=Metric(name="outcome"), op=ComparisonOp.GEQ, bound=-0.5 ) @@ -1324,6 +1325,8 @@ def test_experiment_outcome_constraint_updates(self) -> None: ) # add a scalarized outcome constraint + experiment.add_tracking_metric(Metric(name="oc_m3")) + experiment.add_tracking_metric(Metric(name="oc_m4")) outcome_constraint3 = get_scalarized_outcome_constraint() optimization_config.outcome_constraints = [ outcome_constraint, @@ -1617,12 +1620,10 @@ def test_preference_optimization_config_sqa_roundtrip(self) -> None: loaded_config.objective, MultiObjective ) self.assertEqual( - len(loaded_objective.objectives), - len(multi_objective.objectives), + len(loaded_objective.metric_names), + len(multi_objective.metric_names), ) - loaded_metric_names = { - obj.metric.name for obj in loaded_objective.objectives - } + loaded_metric_names = set(loaded_objective.metric_names) self.assertEqual(loaded_metric_names, {"m1", "m2"}) def test_experiment_objective_threshold_updates(self) -> None: @@ -1643,12 +1644,13 @@ def test_experiment_objective_threshold_updates(self) -> None: objective_threshold, objective_threshold2, ] + experiment.add_tracking_metric(Metric(name="m3")) experiment.optimization_config = optimization_config save_experiment(experiment) self.assertEqual(get_session().query(SQAMetric).count(), 7) - self.assertIsNotNone(optimization_config.objective_thresholds[0].metric.db_id) # add outcome constraint + experiment.add_tracking_metric(Metric(name="outcome")) outcome_constraint2 = OutcomeConstraint( metric=Metric(name="outcome"), op=ComparisonOp.GEQ, bound=-0.5 ) @@ -1665,7 +1667,7 @@ def test_experiment_objective_threshold_updates(self) -> None: optimization_config.outcome_constraints = [] experiment.optimization_config = optimization_config save_experiment(experiment) - self.assertEqual(get_session().query(SQAMetric).count(), 6) + self.assertEqual(get_session().query(SQAMetric).count(), 8) loaded_experiment = load_experiment(experiment.name) self.assertEqual(experiment, loaded_experiment) @@ -1674,7 +1676,7 @@ def test_experiment_objective_threshold_updates(self) -> None: # objective_thresholds optimization_config.objective_thresholds = [] save_experiment(experiment) - self.assertEqual(get_session().query(SQAMetric).count(), 4) + self.assertEqual(get_session().query(SQAMetric).count(), 6) loaded_experiment = load_experiment(experiment.name) self.assertEqual(experiment, loaded_experiment) diff --git a/ax/storage/sqa_store/utils.py b/ax/storage/sqa_store/utils.py index 2986297b2c6..626794dfcdf 100644 --- a/ax/storage/sqa_store/utils.py +++ b/ax/storage/sqa_store/utils.py @@ -80,7 +80,10 @@ def copy_db_ids(source: Any, target: Any, path: list[str] | None = None) -> None raise SQADecodeError(error_message_prefix + "Encountered path of length > 15.") if type(source) is not type(target): - if not issubclass(type(target), type(source)): + if not ( + issubclass(type(target), type(source)) + or issubclass(type(source), type(target)) + ): if source is None and isinstance(target, SearchSpace): warnings.warn( error_message_prefix + "Encountered two objects of different " diff --git a/ax/utils/common/complexity_utils.py b/ax/utils/common/complexity_utils.py index 9c65e0e324d..c43d9f24f97 100644 --- a/ax/utils/common/complexity_utils.py +++ b/ax/utils/common/complexity_utils.py @@ -11,7 +11,6 @@ from ax.adapter.adapter_utils import can_map_to_binary, is_unordered_choice from ax.core.experiment import Experiment -from ax.core.objective import MultiObjective from ax.exceptions.core import OptimizationNotConfiguredError, UserInputError from ax.orchestration.orchestrator import OrchestratorOptions @@ -174,8 +173,8 @@ def summarize_ax_optimization_complexity( num_categorical_6_inf = sum(is_unordered_choice(p, min_choices=6) for p in params) num_parameter_constraints = len(search_space.parameter_constraints) num_objectives = ( - len(optimization_config.objective.objectives) - if isinstance(optimization_config.objective, MultiObjective) + len(optimization_config.objective.metric_names) + if optimization_config.objective.is_multi_objective else 1 ) num_outcome_constraints = len(optimization_config.outcome_constraints) @@ -184,7 +183,9 @@ def summarize_ax_optimization_complexity( # Check if any metrics use merge_multiple_curves uses_merge_multiple_curves = False - all_metrics = list(optimization_config.metrics.values()) + all_metrics = [ + experiment.get_metric(name) for name in optimization_config.metric_names + ] if hasattr(experiment, "tracking_metrics"): all_metrics.extend(experiment.tracking_metrics) for metric in all_metrics: diff --git a/ax/utils/common/string_utils.py b/ax/utils/common/string_utils.py index 0306c6f6d4f..117658d93bc 100644 --- a/ax/utils/common/string_utils.py +++ b/ax/utils/common/string_utils.py @@ -15,10 +15,15 @@ COLON_PLACEHOLDER = "__colon__" PIPE_PLACEHOLDER = "__pipe__" TILDE_PLACEHOLDER = "__tilde__" +SPACE_PLACEHOLDER = "__space__" +HYPHEN_PLACEHOLDER = "__hyphen__" _forbidden_re: re.Pattern[str] = re.compile(r"[\;\[\'\\]") SYMPY_GLOBALS: set[str] = set(dir(sympy)) +# Allow some globals, like oo (infinity) to be used in expressions. +ALLOWED_GLOBALS: set[str] = {"oo"} + def _check_sympy_conflicts(expression: str) -> None: """ @@ -32,7 +37,7 @@ def _check_sympy_conflicts(expression: str) -> None: identifiers = set(re.findall(identifier_pattern, expression)) # Check for conflicts - conflicts = identifiers & SYMPY_GLOBALS + conflicts = (identifiers & SYMPY_GLOBALS) - ALLOWED_GLOBALS if conflicts: conflicts_list = ", ".join(sorted(conflicts)) @@ -60,26 +65,34 @@ def sanitize_name(s: str) -> str: """ if _forbidden_re.search(s) is not None: raise ValueError(f"Expression {s} has forbidden control characters.") - # Replace occurances of "." and "/" when they appear after a valid Python variable - # name and before any alphanumeric character. + # Replace spaces between word characters so metric names with spaces + # (e.g. "CIFAR10 Test Accuracy") become valid Python identifiers. + # Spaces around operators (e.g. "m1 + m2") are not affected because + # operators like +, -, *, are not word characters. + sans_space = re.sub(r"(?<=\w)\s+(?=\w)", SPACE_PLACEHOLDER, s) + # Replace occurrences of "." and "/" when they appear after a valid Python variable + # name and before any alphanumeric character. We use a lookahead (?=...) + # for the trailing character so it is NOT consumed by the match; this + # allows chained separators (e.g. "a:b:c:d") to be fully sanitized in a + # single pass. sans_dots = re.sub( - r"([a-zA-Z_][a-zA-Z0-9_]*)\.([a-zA-Z0-9_])", - rf"\1{DOT_PLACEHOLDER}\2", - s, + r"([a-zA-Z_][a-zA-Z0-9_]*)\.(?=[a-zA-Z0-9_])", + rf"\1{DOT_PLACEHOLDER}", + sans_space, ) sans_slash = re.sub( - r"([a-zA-Z_][a-zA-Z0-9_]*)\/([a-zA-Z0-9_])", - rf"\1{SLASH_PLACEHOLDER}\2", + r"([a-zA-Z_][a-zA-Z0-9_]*)\/(?=[a-zA-Z0-9_])", + rf"\1{SLASH_PLACEHOLDER}", sans_dots, ) sans_colon = re.sub( - r"([a-zA-Z_][a-zA-Z0-9_]*):([a-zA-Z0-9_])", - rf"\1{COLON_PLACEHOLDER}\2", + r"([a-zA-Z_][a-zA-Z0-9_]*):(?=[a-zA-Z0-9_])", + rf"\1{COLON_PLACEHOLDER}", sans_slash, ) sans_pipe = re.sub( - r"([a-zA-Z_][a-zA-Z0-9_]*)\|([a-zA-Z0-9_])", - rf"\1{PIPE_PLACEHOLDER}\2", + r"([a-zA-Z_][a-zA-Z0-9_]*)\|(?=[a-zA-Z0-9_])", + rf"\1{PIPE_PLACEHOLDER}", sans_colon, ) # Replace tilde at the start of a variable name or after alphanumeric characters @@ -88,10 +101,19 @@ def sanitize_name(s: str) -> str: rf"{TILDE_PLACEHOLDER}\1", sans_pipe, ) + # Replace hyphens when they appear between identifier characters + # (e.g. "120s-300s" in metric names like "metric:120s-300s"). + # We require at least 2 characters in the preceding identifier to avoid + # false positives on scientific notation (e.g. "2.3E-4"). + sans_hyphen = re.sub( + r"([a-zA-Z_][a-zA-Z0-9_]+)-(?=[a-zA-Z0-9_])", + rf"\1{HYPHEN_PLACEHOLDER}", + sans_tilde, + ) # Check for conflicts with sympy's global dictionary - _check_sympy_conflicts(sans_tilde) + _check_sympy_conflicts(sans_hyphen) - return sans_tilde + return sans_hyphen def unsanitize_name(s: str) -> str: @@ -103,10 +125,12 @@ def unsanitize_name(s: str) -> str: """ # Unsanitize in the reverse order of sanitization - with_tilde = re.sub(rf"{TILDE_PLACEHOLDER}", "~", s) + with_hyphen = re.sub(rf"{HYPHEN_PLACEHOLDER}", "-", s) + with_tilde = re.sub(rf"{TILDE_PLACEHOLDER}", "~", with_hyphen) with_pipe = re.sub(rf"{PIPE_PLACEHOLDER}", "|", with_tilde) with_colon = re.sub(rf"{COLON_PLACEHOLDER}", ":", with_pipe) with_slash = re.sub(rf"{SLASH_PLACEHOLDER}", "/", with_colon) with_dot = re.sub(rf"{DOT_PLACEHOLDER}", ".", with_slash) + with_space = re.sub(rf"{SPACE_PLACEHOLDER}", " ", with_dot) - return with_dot + return with_space diff --git a/ax/utils/common/sympy.py b/ax/utils/common/sympy.py index ca31d67bf04..f5717fd3ff1 100644 --- a/ax/utils/common/sympy.py +++ b/ax/utils/common/sympy.py @@ -6,10 +6,12 @@ # pyre-strict from ax.exceptions.core import UserInputError -from ax.utils.common.string_utils import sanitize_name +from ax.utils.common.string_utils import sanitize_name, unsanitize_name +from sympy.core.expr import Expr +from sympy.core.numbers import Number from sympy.core.relational import GreaterThan, LessThan from sympy.core.symbol import Symbol -from sympy.core.sympify import sympify +from sympy.core.sympify import sympify, SympifyError def extract_coefficient_dict_from_inequality( @@ -22,7 +24,10 @@ def extract_coefficient_dict_from_inequality( constraints. """ # Parse the constraint string into a SymPy inequality - inequality = sympify(sanitize_name(inequality_str)) + try: + inequality = sympify(sanitize_name(inequality_str)) + except SympifyError: + raise UserInputError(f"Expected an inequality, found {inequality_str}") # Check the SymPy object is a valid inequality if not isinstance(inequality, GreaterThan | LessThan): @@ -38,3 +43,149 @@ def extract_coefficient_dict_from_inequality( return { key: float(value) for key, value in expression.as_coefficients_dict().items() } + + +def parse_objective_expression(expression_str: str) -> Expr | tuple[Expr, ...]: + """Sanitize and sympify an objective expression string. + + For multi-objective expressions (comma-separated), returns a tuple of Expr. + For single-objective expressions, returns a single Expr. + + Args: + expression_str: A string representing the objective expression. + Examples: "m1", "-m1", "2*m1 + m2", "m1, -m2" (multi-objective). + + Returns: + A single SymPy Expr or a tuple of Expr for multi-objective. + + Raises: + UserInputError: If the expression string is empty or invalid. + """ + if len(expression_str) == 0: + raise UserInputError("Objective expression string must not be empty.") + + sanitized = sanitize_name(expression_str) + parsed = sympify(sanitized) + + if isinstance(parsed, tuple): + if any(not isinstance(p, Expr) for p in parsed): + raise UserInputError(f"Invalid objective expression: {expression_str}") + + return parsed + + if not isinstance(parsed, Expr): + raise UserInputError(f"Invalid objective expression: {expression_str}") + + return parsed + + +def extract_metric_names_from_objective_expr(expression: Expr) -> list[str]: + """Extract unsanitized metric names from a SymPy objective expression. + + Uses ``expression.free_symbols`` to find all referenced metrics and + unsanitizes their names (reverting SymPy-safe placeholder substitutions). + + Args: + expression: A SymPy expression representing an objective. + + Returns: + A sorted list of metric name strings (unsanitized). + """ + return sorted(unsanitize_name(str(s)) for s in expression.free_symbols) + + +def extract_metric_weights_from_objective_expr( + expression: Expr, +) -> list[tuple[str, float]]: + """Extract (metric_name, weight) pairs from a SymPy objective expression. + + Handles the following cases: + - Symbol ``x`` -> ``[("x", 1.0)]`` + - Negation ``-x`` -> ``[("x", -1.0)]`` + - Weighted sum ``2*x + y`` -> ``[("x", 2.0), ("y", 1.0)]`` + + Uses ``as_coefficients_dict`` to decompose linear expressions. + + Args: + expression: A SymPy expression representing a (single) objective. + + Returns: + A list of ``(metric_name, weight)`` tuples sorted by metric name. + + Raises: + UserInputError: If the expression contains a constant term (a term + with no metric symbol). + """ + coeff_dict = expression.as_coefficients_dict() + + result: list[tuple[str, float]] = [] + for term, coeff in coeff_dict.items(): + if isinstance(term, Number) or term.is_number: + # Constant offset — skip if zero, error if non-zero + if float(coeff) * float(term) != 0.0: + raise UserInputError( + f"Objective expression '{expression}' contains a constant term. " + "Objective expressions must be purely symbolic." + ) + continue + # term is a Symbol (possibly multiplied by a coefficient already + # absorbed into coeff) + metric_names_in_term = [unsanitize_name(str(s)) for s in term.free_symbols] + if len(metric_names_in_term) != 1: + raise UserInputError( + f"Objective expression '{expression}' contains a non-linear term " + f"'{term}'. Objective expressions must be linear combinations of " + "metrics." + ) + result.append((metric_names_in_term[0], float(coeff))) + + return sorted(result, key=lambda x: x[0]) + + +def _format_number(value: float) -> str: + """Format a float, dropping the trailing '.0' for whole numbers.""" + return str(int(value)) if value == int(value) else str(value) + + +def build_constraint_expression_str( + metric_weights: list[tuple[str, float]], + op: str, + bound: float, + relative: bool, +) -> str: + """Reconstruct a constraint expression string from components. + + Builds a human-readable inequality string from metric weights, comparison + operator, bound, and relativity. + + Args: + metric_weights: List of ``(metric_name, weight)`` tuples. + op: The comparison operator string, either ``">="`` or ``"<="``. + bound: The numeric bound. + relative: Indicates whether the constraint is relative to the status quo. + + Returns: + A constraint expression string, e.g. ``"qps >= 700"`` or + ``"loss <= 0.5 * baseline"``. + """ + parts: list[str] = [] + for name, weight in metric_weights: + if weight == 1.0: + parts.append(name) + elif weight == -1.0: + parts.append(f"-{name}") + else: + parts.append(f"{_format_number(weight)}*{name}") + + lhs = " + ".join(parts).replace(" + -", " - ") + + if relative: + multiplier = 1 + bound / 100 + if multiplier == 1: + bound_str = "baseline" + else: + bound_str = f"{_format_number(multiplier)} * baseline" + return f"{lhs} {op} {bound_str}" + + bound_str = _format_number(bound) + return f"{lhs} {op} {bound_str}" diff --git a/ax/utils/common/tests/test_string_utils.py b/ax/utils/common/tests/test_string_utils.py index fe557f810f7..fea6b749e78 100644 --- a/ax/utils/common/tests/test_string_utils.py +++ b/ax/utils/common/tests/test_string_utils.py @@ -63,5 +63,8 @@ def test_sanitize_name_sympy_conflicts(self) -> None: # Test that E in scientific notation is not flagged (it's not an identifier) self.assertEqual(sanitize_name("1E5 + 2.3E-4"), "1E5 + 2.3E-4") - # Test that math functions are not flagged - self.assertEqual(sanitize_name("sin(1) + cos(2)"), "sin(1) + cos(2)") + # Test that parens and equals are preserved (not sanitized). + self.assertEqual( + sanitize_name("sin(1) + cos(2)"), + "sin(1) + cos(2)", + ) diff --git a/ax/utils/testing/core_stubs.py b/ax/utils/testing/core_stubs.py index a2dcee5d5dd..7c143d45463 100644 --- a/ax/utils/testing/core_stubs.py +++ b/ax/utils/testing/core_stubs.py @@ -17,7 +17,7 @@ from logging import Logger from math import prod from pathlib import Path -from typing import Any, cast, Self +from typing import Any, Self import numpy as np import pandas as pd @@ -173,7 +173,7 @@ def get_experiment_with_map_data_type() -> Experiment: optimization_config=get_map_optimization_config(), status_quo=get_status_quo(), description="test description", - tracking_metrics=[MapMetric(name="tracking")], + tracking_metrics=[MapMetric(name="m1"), MapMetric(name="tracking")], is_test=True, ) experiment._properties = {"owners": [DEFAULT_USER]} @@ -262,15 +262,44 @@ def get_experiment_with_custom_runner_and_metric( outcome_constraint=has_outcome_constraint, relative=False ) + # Build tracking_metrics with actual metric instances so the + # Experiment owns the real metric types. + experiment_tracking_metrics: list[Metric] = [ + CustomTestMetric(name="custom_test_metric", test_attribute="test") + ] + if multi_objective: + experiment_tracking_metrics.extend( + [ + CustomTestMetric(name="m1", test_attribute="test"), + CustomTestMetric(name="m3", test_attribute="test"), + ] + ) + if has_outcome_constraint: + experiment_tracking_metrics.append( + Metric(name="m2"), + ) + elif scalarized_objective: + experiment_tracking_metrics.extend( + [ + CustomTestMetric(name="m1", test_attribute="test"), + CustomTestMetric(name="m3", test_attribute="test"), + ] + ) + if has_outcome_constraint: + experiment_tracking_metrics.extend( + [ + CustomTestMetric(name="oc_m3", test_attribute="test"), + CustomTestMetric(name="oc_m4", test_attribute="test"), + ] + ) + # Create experiment with custom runner and metric experiment = Experiment( name="test", search_space=get_search_space(constrain_search_space=constrain_search_space), optimization_config=optimization_config, description="test description", - tracking_metrics=[ - CustomTestMetric(name="custom_test_metric", test_attribute="test") - ], + tracking_metrics=experiment_tracking_metrics, runner=CustomTestRunner(test_attribute="test"), is_test=True, ) @@ -348,6 +377,15 @@ def get_branin_experiment( else None ) + # Build tracking_metrics with actual BraninMetric instances so the + # Experiment owns the real metric types (the optimization config only + # stores metric names, not Metric objects). + branin_tracking: list[Metric] = [get_branin_metric(name="branin")] + if with_relative_constraint: + branin_tracking.append(get_branin_metric(name="branin_d")) + if with_absolute_constraint: + branin_tracking.append(get_branin_metric(name="branin_e")) + exp = Experiment( name="branin_test_experiment" if named else None, search_space=search_space, @@ -360,6 +398,7 @@ def get_branin_experiment( if has_optimization_config or with_relative_constraint else None ), + tracking_metrics=branin_tracking, runner=SyntheticRunner(), is_test=True, status_quo=status_quo, @@ -492,7 +531,7 @@ def get_branin_experiment_with_timestamp_map_metric( num_objectives = 2 bounds = bounds or [99.0 for _ in range(num_objectives)] if has_objective_thresholds: - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] | None = [ ObjectiveThreshold( metric=local_get_map_metric(name=f"branin_map_{m}"), bound=bound, @@ -516,6 +555,16 @@ def get_branin_experiment_with_timestamp_map_metric( Objective(metric=local_get_map_metric(f"branin_map_{m}")) for m in range(num_objectives) ] + # Add objective metrics so the Experiment owns the real metric types. + # pyre-ignore[6]: Covariance issue with list[T] + tracking_metrics.extend( + local_get_map_metric(f"branin_map_{m}") for m in range(num_objectives) + ) + if with_outcome_constraint: + tracking_metrics.append( + # pyre-ignore[6]: Covariance issue with list[T] + BraninMetric(name="branin_constraint", param_names=["x1", "x2"]) + ) optimization_config = MultiObjectiveOptimizationConfig( objective=MultiObjective(objectives=objectives), objective_thresholds=objective_thresholds, @@ -539,6 +588,12 @@ def get_branin_experiment_with_timestamp_map_metric( ) tracking_metrics = [tracking_metric] + # Add objective metric so the Experiment owns the real metric type. + tracking_metrics.append(local_get_map_metric(name="branin_map")) + if with_outcome_constraint: + tracking_metrics.append( + BraninMetric(name="branin_constraint", param_names=["x1", "x2"]) + ) exp = Experiment( name=experiment_name, @@ -546,7 +601,8 @@ def get_branin_experiment_with_timestamp_map_metric( with_choice_parameter=with_choice_parameter ), optimization_config=optimization_config, - tracking_metrics=cast(list[Metric], tracking_metrics), + # pyre-ignore[6] Covariance issue with list[T] + tracking_metrics=tracking_metrics, runner=SyntheticRunner(), ) exp._properties = {"owners": [DEFAULT_USER]} @@ -619,7 +675,9 @@ def get_test_map_data_experiment( def get_multi_type_experiment( add_trial_type: bool = True, add_trials: bool = False, num_arms: int = 10 ) -> MultiTypeExperiment: - oc = OptimizationConfig(Objective(BraninMetric("m1", ["x1", "x2"]), minimize=True)) + oc = OptimizationConfig( + Objective(metric=BraninMetric("m1", ["x1", "x2"]), minimize=True) + ) experiment = MultiTypeExperiment( name="test_exp", search_space=get_branin_search_space(), @@ -627,6 +685,7 @@ def get_multi_type_experiment( default_runner=SyntheticRunner(dummy_metadata="dummy1"), optimization_config=oc, status_quo=Arm(parameters={"x1": 0.0, "x2": 0.0}), + tracking_metrics=[BraninMetric("m1", ["x1", "x2"])], ) experiment._properties = {"owners": [DEFAULT_USER]} experiment.add_trial_type( @@ -669,7 +728,10 @@ def get_factorial_experiment( ), runner=SyntheticRunner(), is_test=True, - tracking_metrics=[get_factorial_metric("secondary_metric")], + tracking_metrics=[ + get_factorial_metric(), + get_factorial_metric("secondary_metric"), + ], ) exp._properties = {"owners": [DEFAULT_USER]} @@ -751,7 +813,11 @@ def get_experiment_with_multi_objective() -> Experiment: optimization_config=optimization_config, description="test experiment with multi objective", runner=SyntheticRunner(), - tracking_metrics=[Metric(name="tracking")], + tracking_metrics=[ + Metric(name="m1", lower_is_better=False), + Metric(name="m3", lower_is_better=True), + Metric(name="tracking"), + ], is_test=True, ) exp._properties = {"owners": [DEFAULT_USER]} @@ -809,6 +875,20 @@ def get_branin_experiment_with_multi_objective( if has_optimization_config else None ) + # Build tracking_metrics with actual BraninMetric instances so the + # Experiment owns the real metric types (the optimization config only + # stores metric names, not Metric objects). + branin_tracking: list[Metric] = [ + get_branin_metric(name="branin_a"), + get_branin_metric(name="branin_b"), + ] + if num_objectives == 3: + branin_tracking.append(get_branin_metric(name="branin_c")) + if with_relative_constraint: + branin_tracking.append(get_branin_metric(name="branin_d")) + if with_absolute_constraint: + branin_tracking.append(get_branin_metric(name="branin_e")) + exp = Experiment( name="branin_test_experiment", search_space=get_branin_search_space( @@ -818,6 +898,7 @@ def get_branin_experiment_with_multi_objective( with_derived_parameter=with_derived_parameter, ), optimization_config=optimization_config, + tracking_metrics=branin_tracking, runner=SyntheticRunner(), is_test=True, ) @@ -875,6 +956,16 @@ def get_branin_experiment_with_multi_objective( def get_branin_with_multi_task(with_multi_objective: bool = False) -> Experiment: + # Build tracking_metrics with actual BraninMetric instances so the + # Experiment owns the real metric types. + if with_multi_objective: + branin_tracking: list[Metric] = [ + get_branin_metric(name="branin_a"), + get_branin_metric(name="branin_b"), + ] + else: + branin_tracking = [get_branin_metric(name="branin")] + exp = Experiment( name="branin_test_experiment", search_space=get_branin_search_space(), @@ -885,6 +976,7 @@ def get_branin_with_multi_task(with_multi_objective: bool = False) -> Experiment if with_multi_objective else get_branin_optimization_config() ), + tracking_metrics=branin_tracking, # pyre-ignore[6] covariance runner=SyntheticRunner(), is_test=True, ) @@ -1015,6 +1107,7 @@ def get_experiment_with_observations( signature_override="m2_sig" if signature_suffix else None, ), ] + tracking_metrics_from_opt_config = list(metrics) if scalarized: optimization_config = OptimizationConfig( objective=ScalarizedObjective(metrics) @@ -1022,6 +1115,16 @@ def get_experiment_with_observations( if constrained: raise NotImplementedError else: + constraint_metric = ( + Metric( + name="m3", + signature_override="m3_sig" if signature_suffix else None, + ) + if constrained + else None + ) + if constraint_metric is not None: + tracking_metrics_from_opt_config.append(constraint_metric) optimization_config = MultiObjectiveOptimizationConfig( objective=MultiObjective( objectives=[Objective(metric=metric) for metric in metrics] @@ -1042,12 +1145,7 @@ def get_experiment_with_observations( outcome_constraints=( [ OutcomeConstraint( - metric=Metric( - name="m3", - signature_override="m3_sig" - if signature_suffix - else None, - ), + metric=constraint_metric, op=ComparisonOp.GEQ, bound=0.0, relative=False, @@ -1060,19 +1158,23 @@ def get_experiment_with_observations( elif optimization_config is None: if scalarized: raise NotImplementedError + obj_metric = Metric( + name="m1", + signature_override="m1_sig" if signature_suffix else None, + ) objective = Objective( - metric=Metric( - name="m1", - signature_override="m1_sig" if signature_suffix else None, - ), + metric=obj_metric, minimize=minimize, ) + tracking_metrics_from_opt_config = [obj_metric] if constrained: + constraint_metric = Metric( + name="m2", + signature_override="m2_sig" if signature_suffix else None, + ) + tracking_metrics_from_opt_config.append(constraint_metric) constraint = OutcomeConstraint( - metric=Metric( - name="m2", - signature_override="m2_sig" if signature_suffix else None, - ), + metric=constraint_metric, op=ComparisonOp.GEQ, bound=0.0, relative=False, @@ -1082,23 +1184,28 @@ def get_experiment_with_observations( ) else: optimization_config = OptimizationConfig(objective=objective) + else: + tracking_metrics_from_opt_config = [] search_space = search_space or get_search_space_for_range_values(min=0.0, max=1.0) + # Collect tracking metrics: include metrics from the optimization config + # (so their signature_override and lower_is_better are preserved on the + # experiment, rather than being replaced by plain Metric placeholders) + # plus any additional tracking metrics requested by the caller. + all_tracking_metrics = list(tracking_metrics_from_opt_config) + if with_tracking_metrics: + all_tracking_metrics.append( + Metric( + name=f"m{len(observations[0])}", + lower_is_better=False, + signature_override=f"m{len(observations[0])}_sig" + if signature_suffix + else None, + ) + ) exp = Experiment( search_space=search_space, optimization_config=optimization_config, - tracking_metrics=( - [ - Metric( - name=f"m{len(observations[0])}", - lower_is_better=False, - signature_override=f"m{len(observations[0])}_sig" - if signature_suffix - else None, - ) - ] - if with_tracking_metrics - else None - ), + tracking_metrics=all_tracking_metrics or None, runner=SyntheticRunner(), is_test=True, status_quo=status_quo, @@ -1187,6 +1294,9 @@ def get_high_dimensional_branin_experiment( name="high_dimensional_branin_experiment", search_space=search_space, optimization_config=optimization_config, + tracking_metrics=[ + BraninMetric(name="objective", param_names=["x19", "x44"]), + ], runner=SyntheticRunner(), status_quo=Arm(sq_parameters) if with_status_quo else None, ) @@ -2194,7 +2304,7 @@ def get_multi_objective_optimization_config( outcome_constraints = ( [get_outcome_constraint(relative=relative)] if outcome_constraint else [] ) - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] = [ get_objective_threshold(metric_name="m1"), get_objective_threshold(metric_name="m3", comparison_op=ComparisonOp.LEQ), ] @@ -2252,7 +2362,7 @@ def get_branin_multi_objective_optimization_config( _validate_num_outcomes(num_outcomes=num_objectives + num_constraint_outcomes) # minimum Branin value is 0.397887 if has_objective_thresholds: - objective_thresholds = [ + objective_thresholds: list[OutcomeConstraint] = [ ObjectiveThreshold( metric=get_branin_metric(name="branin_a"), bound=10.0, diff --git a/ax/utils/testing/utils.py b/ax/utils/testing/utils.py index daa2feb4702..e1d1d3de151 100644 --- a/ax/utils/testing/utils.py +++ b/ax/utils/testing/utils.py @@ -72,7 +72,7 @@ def run_trials_with_gs( "trial_index": i, "metric_signature": m, } - for m in none_throws(experiment.optimization_config).metrics + for m in none_throws(experiment.optimization_config).metric_names for arm in trial.arms ] ) diff --git a/tutorials/external_generation_node/external_generation_node.ipynb b/tutorials/external_generation_node/external_generation_node.ipynb index 9138cf4d132..a05e50bf245 100644 --- a/tutorials/external_generation_node/external_generation_node.ipynb +++ b/tutorials/external_generation_node/external_generation_node.ipynb @@ -122,7 +122,7 @@ " \"\"\"\n", " search_space = experiment.search_space\n", " parameter_names = list(search_space.parameters.keys())\n", - " metric_names = list(experiment.optimization_config.metrics.keys())\n", + " metric_names = experiment.optimization_config.metric_names\n", " if any(\n", " not isinstance(p, RangeParameter) for p in search_space.parameters.values()\n", " ):\n",