From 4aeb120d38ce19ce52f0911e9ca1ead84a4a67e9 Mon Sep 17 00:00:00 2001 From: Sait Cakmak Date: Thu, 22 Oct 2026 09:30:25 -0700 Subject: [PATCH] Preserve multi-metric outcome constraints through candidate generation Summary: Fixes https://github.com/facebook/Ax/issues/5309. What Changed: - Extract all signed coefficients from `OutcomeConstraint.metric_weights`, without depending on the deprecated `ScalarizedOutcomeConstraint` class. - Validate every fitted constraint metric and compute relative baseline values and diagnostics from every component. - Exclude scalarized constraints from automatic single-metric winsorization and avoid nonlinear `BilogY` transformations of their component metrics, including components that also have single-metric constraints. - Reject relative multi-metric expressions in the Client API parser, consistent with the existing documented restriction. Lower-level mixed-sign relative parsing remains a separate issue. - Add regression coverage for expression constraints, both inequality directions, signed coefficients, canonical signatures, relative baseline handling, missing component data, and end-to-end Client generation. The Client reproduction produced `A=[[0,1,0]]` before the fix. The regression now verifies both appropriately standardized metric coefficients and the transformed bound. Differential Revision: D126924134 --- ax/adapter/adapter_utils.py | 16 +-- ax/adapter/tests/test_utils.py | 53 +++++++ ax/adapter/torch.py | 8 +- ax/adapter/transforms/bilog_y.py | 12 +- ax/adapter/transforms/derelativize.py | 36 ++--- ax/adapter/transforms/tests/test_bilog_y.py | 35 +++++ .../tests/test_derelativize_transform.py | 132 +++++++++++++++++- .../tests/test_winsorize_transform.py | 53 ++++++- ax/adapter/transforms/winsorize.py | 7 +- ax/api/tests/test_client.py | 52 +++++++ ax/api/utils/instantiation/from_string.py | 7 + .../instantiation/tests/test_from_string.py | 34 +++++ 12 files changed, 385 insertions(+), 60 deletions(-) diff --git a/ax/adapter/adapter_utils.py b/ax/adapter/adapter_utils.py index 8743e7b3103..a65aceeb12d 100644 --- a/ax/adapter/adapter_utils.py +++ b/ax/adapter/adapter_utils.py @@ -34,11 +34,7 @@ OptimizationConfig, TRefPoint, ) -from ax.core.outcome_constraint import ( - ComparisonOp, - OutcomeConstraint, - ScalarizedOutcomeConstraint, -) +from ax.core.outcome_constraint import ComparisonOp, OutcomeConstraint from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter from ax.core.parameter_constraint import ParameterConstraint from ax.core.search_space import SearchSpace, SearchSpaceDigest @@ -406,13 +402,9 @@ def extract_outcome_constraints( b = np.zeros((len(outcome_constraints), 1)) for i, c in enumerate(outcome_constraints): s = 1 if c.op == ComparisonOp.LEQ else -1 - if isinstance(c, ScalarizedOutcomeConstraint): - for c_metric_sig, c_weight in c.metric_weights: - j = outcomes.index(c_metric_sig) - A[i, j] = s * c_weight - else: - j = outcomes.index(c.metric_signatures[0]) - A[i, j] = s + for c_metric_sig, c_weight in c.metric_weights: + j = outcomes.index(c_metric_sig) + A[i, j] = s * c_weight b[i, 0] = s * c.bound return (A, b) diff --git a/ax/adapter/tests/test_utils.py b/ax/adapter/tests/test_utils.py index dd78ea31343..8979059442a 100644 --- a/ax/adapter/tests/test_utils.py +++ b/ax/adapter/tests/test_utils.py @@ -16,12 +16,14 @@ pending_observations_as_array_list, ) from ax.adapter.registry import Generators +from ax.adapter.torch import validate_transformed_optimization_config from ax.core.arm import Arm from ax.core.evaluations_to_data import raw_evaluations_to_data from ax.core.generator_run import GeneratorRun from ax.core.metric import Metric from ax.core.objective import MultiObjective, Objective from ax.core.observation import ObservationData, ObservationFeatures +from ax.core.optimization_config import OptimizationConfig from ax.core.outcome_constraint import ( ObjectiveThreshold, OutcomeConstraint, @@ -29,6 +31,7 @@ ) from ax.core.types import ComparisonOp from ax.core.utils import get_pending_observation_features +from ax.exceptions.core import DataRequiredError from ax.utils.common.constants import Keys from ax.utils.common.testutils import TestCase from ax.utils.testing.core_stubs import ( @@ -116,6 +119,56 @@ def test_extract_outcome_constraints(self) -> None: self.assertEqual(res[1][0][0], 0) self.assertEqual(res[1][1][0], -1) + def test_expression_outcome_constraints(self) -> None: + """Preserve every coefficient and resolve names to fitted signatures.""" + mapping = {"a": "a_sig", "b": "b_sig"} + constraints = [ + OutcomeConstraint(expression=expression, metric_name_to_signature=mapping) + for expression in [ + "a + 2 * b <= 1.0", + "a + 2 * b >= 1.0", + "a - 2 * b >= -3.0", + "2 * a <= 4.0", + ] + ] + A, b = none_throws( + extract_outcome_constraints( + outcome_constraints=constraints, outcomes=["b_sig", "obj", "a_sig"] + ) + ) + np.testing.assert_allclose(A, [[2, 0, 1], [-2, 0, -1], [2, 0, -1], [0, 0, 1]]) + np.testing.assert_allclose(b, [[1], [-1], [3], [2]]) + + def test_expression_constraint_fitted_metrics(self) -> None: + """Validate every component, not just the first metric in an expression.""" + config = OptimizationConfig( + objective=Objective( + expression="obj", metric_name_to_signature={"obj": "obj"} + ), + outcome_constraints=[ + OutcomeConstraint( + expression="a + 2 * b <= 1.0", + metric_name_to_signature={"a": "a_sig", "b": "b_sig"}, + ) + ], + ) + validate_transformed_optimization_config( + optimization_config=config, outcomes=["obj", "a_sig", "b_sig"] + ) + for missing in ["a_sig", "b_sig"]: + with ( + self.subTest(missing=missing), + self.assertRaisesRegex( + DataRequiredError, f"{missing} not found in fitted data" + ), + ): + validate_transformed_optimization_config( + optimization_config=config, + outcomes=[ + sig for sig in ["obj", "a_sig", "b_sig"] if sig != missing + ], + ) + def test_extract_objective_thresholds(self) -> None: outcomes = ["m1", "m2", "m3", "m4"] objective = MultiObjective( diff --git a/ax/adapter/torch.py b/ax/adapter/torch.py index b8c33ae09be..a95493901cc 100644 --- a/ax/adapter/torch.py +++ b/ax/adapter/torch.py @@ -72,11 +72,7 @@ OptimizationConfig, PreferenceOptimizationConfig, ) -from ax.core.outcome_constraint import ( - ComparisonOp, - ObjectiveThreshold, - ScalarizedOutcomeConstraint, -) +from ax.core.outcome_constraint import ComparisonOp, ObjectiveThreshold from ax.core.search_space import SearchSpace from ax.core.types import TCandidateMetadata, TModelPredictArm from ax.exceptions.core import DataRequiredError, UnsupportedError, UserInputError @@ -1334,7 +1330,7 @@ def validate_transformed_optimization_config( "either a `Relativize` or `Derelativize` transform was expected to be" " applied but was not." ) - if isinstance(c, ScalarizedOutcomeConstraint): + if len(c.metric_weights) > 1: for c_sig in c.metric_signatures: if c_sig not in outcomes: raise DataRequiredError( diff --git a/ax/adapter/transforms/bilog_y.py b/ax/adapter/transforms/bilog_y.py index 9ad2274940e..24c866cad2a 100644 --- a/ax/adapter/transforms/bilog_y.py +++ b/ax/adapter/transforms/bilog_y.py @@ -66,12 +66,20 @@ def __init__( config=config, ) if adapter is not None and adapter._optimization_config is not None: + constraints = adapter._optimization_config.outcome_constraints + scalarized_metrics = { + sig + for oc in constraints + if len(oc.metric_weights) > 1 + for sig in oc.metric_signatures + } # TODO @deriksson: Add support for relative outcome constraints self.metric_to_bound: dict[str, float] = { sig: oc.bound - for oc in adapter._optimization_config.outcome_constraints - if not oc.relative + for oc in constraints + if not oc.relative and len(oc.metric_weights) == 1 for sig in oc.metric_signatures + if sig not in scalarized_metrics } else: self.metric_to_bound = {} diff --git a/ax/adapter/transforms/derelativize.py b/ax/adapter/transforms/derelativize.py index 1357a1f8ee2..b6d5ef9ee71 100644 --- a/ax/adapter/transforms/derelativize.py +++ b/ax/adapter/transforms/derelativize.py @@ -20,7 +20,7 @@ MultiObjectiveOptimizationConfig, OptimizationConfig, ) -from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint +from ax.core.outcome_constraint import OutcomeConstraint from ax.core.types import ComparisonOp, TModelCov, TModelMean from ax.exceptions.core import DataRequiredError from ax.utils.common.logger import get_logger @@ -174,26 +174,14 @@ def _compute_sq_val( 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(c.metric_weights) > 1: + missing_metrics = {sig for sig, _ in c.metric_weights if sig 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) - ] - ) - ) + return float(np.sum([weight * f[sig][0] for sig, weight in c.metric_weights])) else: metric = adapter._experiment.get_metrics(c.metric_names)[0] if metric.signature in f: @@ -215,18 +203,14 @@ def _warn_if_raw_sq_is_out_of_CI( by more than 1.96 standard deviation from the predictions. """ relative_metrics = { - oc.metric_signatures[0] + sig for oc in optimization_config.all_constraints - if oc.relative and not isinstance(oc, ScalarizedOutcomeConstraint) - }.union( - { - sig - for oc in optimization_config.all_constraints - if oc.relative and isinstance(oc, ScalarizedOutcomeConstraint) - for sig in oc.metric_signatures - } - ) + if oc.relative + for sig in oc.metric_signatures + } for metric_signature in relative_metrics: + if metric_signature not in raw_f or metric_signature not in pred_f: + continue raw_obs = raw_f[metric_signature][0] pred_mean = pred_f[metric_signature][0] pred_std = np.sqrt(pred_cov[metric_signature][metric_signature][0]) diff --git a/ax/adapter/transforms/tests/test_bilog_y.py b/ax/adapter/transforms/tests/test_bilog_y.py index d4cc1279cf0..8e6bff5dd40 100644 --- a/ax/adapter/transforms/tests/test_bilog_y.py +++ b/ax/adapter/transforms/tests/test_bilog_y.py @@ -10,13 +10,18 @@ from copy import deepcopy from functools import partial from itertools import product +from unittest.mock import Mock import numpy as np from ax.adapter.base import Adapter, DataLoaderConfig from ax.adapter.data_utils import extract_experiment_data from ax.adapter.transforms.bilog_y import bilog_transform, BilogY, inv_bilog_transform from ax.adapter.transforms.log_y import match_ci_width +from ax.core.objective import Objective +from ax.core.observation import ObservationData from ax.core.observation_utils import observations_from_data +from ax.core.optimization_config import OptimizationConfig +from ax.core.outcome_constraint import OutcomeConstraint from ax.generators.base import Generator from ax.utils.common.testutils import TestCase from ax.utils.testing.core_stubs import get_branin_experiment @@ -58,6 +63,36 @@ def test_Init(self) -> None: ) self.assertEqual(t.metric_to_bound, {}) + def test_scalarized_constraint_metrics_are_not_transformed(self) -> None: + """Nonlinear component transforms cannot preserve a linear combination.""" + mapping = {"a": "a_sig", "b": "b_sig", "c": "c_sig"} + config = OptimizationConfig( + objective=Objective( + expression="obj", metric_name_to_signature={"obj": "obj"} + ), + outcome_constraints=[ + OutcomeConstraint(expression=expr, metric_name_to_signature=mapping) + for expr in ["a + 2 * b <= 5", "a <= 2", "c <= 1"] + ], + ) + transform = BilogY(adapter=Mock(spec=Adapter, _optimization_config=config)) + self.assertEqual(transform.metric_to_bound, {"c_sig": 1.0}) + observation = ObservationData( + metric_signatures=["a_sig", "b_sig", "c_sig"], + means=np.array([3.0, 4.0, 5.0]), + covariance=np.eye(3), + ) + transformed = transform._transform_observation_data( + observation_data=[deepcopy(observation)] + )[0] + np.testing.assert_allclose(transformed.means[:2], observation.means[:2]) + np.testing.assert_allclose( + transformed.covariance[:2, :2], observation.covariance[:2, :2] + ) + self.assertAlmostEqual( + transformed.means[2], float(bilog_transform(y=np.array(5.0), bound=1)) + ) + def test_Bilog(self) -> None: self.assertAlmostEqual( float(bilog_transform(y=np.array(7.3), bound=3)), 4.667706820558076 diff --git a/ax/adapter/transforms/tests/test_derelativize_transform.py b/ax/adapter/transforms/tests/test_derelativize_transform.py index e756d8a31ec..3e5cd5b7420 100644 --- a/ax/adapter/transforms/tests/test_derelativize_transform.py +++ b/ax/adapter/transforms/tests/test_derelativize_transform.py @@ -14,7 +14,11 @@ import numpy as np from ax.adapter.base import Adapter from ax.adapter.torch import TorchAdapter -from ax.adapter.transforms.derelativize import Derelativize, logger +from ax.adapter.transforms.derelativize import ( + _warn_if_raw_sq_is_out_of_CI, + Derelativize, + logger, +) from ax.core.arm import Arm from ax.core.experiment import Experiment from ax.core.metric import Metric @@ -300,7 +304,7 @@ def _test_DerelativizeTransform( ), self.assertRaisesRegex( DataRequiredError, - "Status-quo metric value not yet available for metric\\(s\\) ", + "Status-quo metric value not yet available for metric ", ), ): t2.transform_optimization_config(deepcopy(oc_scalarized_only), g, None) @@ -316,6 +320,130 @@ def _test_DerelativizeTransform( with self.assertRaises(ValueError): t.transform_optimization_config(deepcopy(oc), None, None) + def test_expression_relative_constraint(self) -> None: + """Derelativize against the same weighted outcome at the status quo.""" + for sign in [1.0, -1.0]: + experiment = get_experiment_with_observations( + observations=[[sign, 2 * sign]], + sems=[[1.0, 1.0]], + parameterizations=[{"x": 2.0, "y": 10.0}], + ).clone_with( + name="relative_constraint_test", + status_quo=Arm(parameters={"x": 2.0, "y": 10.0}, name="0_0"), + ) + adapter = Adapter(experiment=experiment, generator=Generator()) + for use_raw in [False, True]: + with self.subTest(sign=sign, use_raw=use_raw): + constraint = OutcomeConstraint( + expression="m1 + 2 * m2 >= 0.9 * baseline", + metric_name_to_signature={"m1": "m1", "m2": "m2"}, + ) + config = OptimizationConfig( + objective=Objective( + expression="obj", metric_name_to_signature={"obj": "obj"} + ), + outcome_constraints=[constraint], + ) + transform = Derelativize( + search_space=None, config={"use_raw_status_quo": use_raw} + ) + with patch.object( + adapter, + "predict", + return_value=( + {"m1": [3 * sign], "m2": [5 * sign]}, + {"m1": {"m1": [1.0]}, "m2": {"m2": [4.0]}}, + ), + ) as predict: + result = transform.transform_optimization_config( + optimization_config=config, adapter=adapter + ).outcome_constraints[0] + sq_val = (5 if use_raw else 13) * sign + self.assertAlmostEqual(result.bound, sq_val - 0.1 * abs(sq_val)) + self.assertEqual(result.metric_weights, constraint.metric_weights) + self.assertEqual(result.op, constraint.op) + self.assertFalse(result.relative) + self.assertEqual(predict.call_count, int(not use_raw)) + + def test_expression_relative_constraint_missing_metric(self) -> None: + """Require status quo data for every component of a relative expression.""" + experiment = get_experiment_with_observations( + observations=[[1.0, 2.0]], + sems=[[1.0, 1.0]], + parameterizations=[{"x": 2.0, "y": 10.0}], + ).clone_with( + name="relative_constraint_test", + status_quo=Arm(parameters={"x": 2.0, "y": 10.0}, name="0_0"), + ) + adapter = Adapter(experiment=experiment, generator=Generator()) + config = OptimizationConfig( + objective=Objective( + expression="obj", metric_name_to_signature={"obj": "obj"} + ), + outcome_constraints=[ + OutcomeConstraint( + expression="m1 + 2 * m2 >= 0.9 * baseline", + metric_name_to_signature={"m1": "m1", "m2": "m2"}, + ) + ], + ) + transform = Derelativize(search_space=None, config={"use_raw_status_quo": True}) + with ( + patch( + f"{Derelativize.__module__}.unwrap_observation_data", + return_value=({"m1": [1.0]}, {}), + ), + self.assertRaisesRegex(DataRequiredError, "m2"), + ): + transform.transform_optimization_config( + optimization_config=config, adapter=adapter + ) + + def test_expression_relative_constraint_warnings(self) -> None: + """Check both canonical metric signatures when warning about the baseline.""" + config = OptimizationConfig( + objective=Objective( + expression="obj", metric_name_to_signature={"obj": "obj"} + ), + outcome_constraints=[ + OutcomeConstraint( + expression="a + 2 * b <= 0.9 * baseline", + metric_name_to_signature={"a": "a_sig", "b": "b_sig"}, + ) + ], + ) + for signature in ["a_sig", "b_sig"]: + with ( + self.subTest(signature=signature), + self.assertLogs(logger=logger) as logs, + ): + _warn_if_raw_sq_is_out_of_CI( + optimization_config=config, + raw_f={"a_sig": [0.0], "b_sig": [0.0]}, + pred_f={ + sig: [10.0 if sig == signature else 0.0] + for sig in ["a_sig", "b_sig"] + }, + pred_cov={"a_sig": {"a_sig": [1.0]}, "b_sig": {"b_sig": [1.0]}}, + ) + self.assertEqual(len(logs.records), 1) + self.assertIn(signature, logs.records[0].getMessage()) + for missing_raw in [False, True]: + with ( + self.subTest(missing_raw=missing_raw), + self.assertNoLogs(logger=logger), + ): + _warn_if_raw_sq_is_out_of_CI( + optimization_config=config, + raw_f={"a_sig": [0.0]} + if missing_raw + else {"a_sig": [0.0], "b_sig": [0.0]}, + pred_f={"a_sig": [0.0], "b_sig": [0.0]} + if missing_raw + else {"a_sig": [0.0]}, + pred_cov={"a_sig": {"a_sig": [1.0]}, "b_sig": {"b_sig": [1.0]}}, + ) + def test_errors(self) -> None: t = Derelativize(search_space=None) oc = OptimizationConfig( diff --git a/ax/adapter/transforms/tests/test_winsorize_transform.py b/ax/adapter/transforms/tests/test_winsorize_transform.py index e67dff8b1c2..93dd62ad137 100644 --- a/ax/adapter/transforms/tests/test_winsorize_transform.py +++ b/ax/adapter/transforms/tests/test_winsorize_transform.py @@ -19,6 +19,7 @@ _get_auto_winsorization_cutoffs_outcome_constraint, _get_auto_winsorization_cutoffs_single_objective, _get_tukey_cutoffs, + _obtain_cutoffs_from_outcome_constraints, AUTO_WINS_QUANTILE, Winsorize, ) @@ -526,6 +527,41 @@ 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"], (-3.5, INF)) # 1 - 1.5 * 3 + def test_expression_outcome_constraints(self) -> None: + """Do not apply single-metric winsorization to a weighted outcome.""" + mapping = {"a": "a_sig", "b": "b_sig"} + config = OptimizationConfig( + objective=Objective( + expression="obj", metric_name_to_signature={"obj": "obj"} + ), + outcome_constraints=[ + OutcomeConstraint( + expression="a + 2 * b <= 1.0", metric_name_to_signature=mapping + ) + ], + ) + for signature in mapping.values(): + with ( + self.subTest(signature=signature), + self.assertWarnsRegex(AxOptimizationWarning, f"manually.*{signature}"), + ): + cutoffs = _obtain_cutoffs_from_outcome_constraints( + optimization_config=config, + metric_signature=signature, + metric_values=[0.0, 1.0, 2.0], + ) + self.assertEqual(cutoffs, (-INF, INF)) + config.outcome_constraints.append( + OutcomeConstraint(expression="a <= 1.0", metric_name_to_signature=mapping) + ) + with self.assertWarnsRegex(AxOptimizationWarning, "manually.*a_sig"): + cutoffs = _obtain_cutoffs_from_outcome_constraints( + optimization_config=config, + metric_signature="a_sig", + metric_values=[0.0, 1.0, 2.0], + ) + self.assertEqual(cutoffs, (-INF, 5.0)) + def test_relative_constraints(self) -> None: # Adapter with in-design status quo search_space = SearchSpace( @@ -615,14 +651,17 @@ def test_relative_constraints(self) -> None: t.cutoffs, {"a": (-INF, INF), "b": (-INF, INF), "c": (-3.25, INF)} ) # Winsorizes with `derelativize_with_raw_status_quo`. - t = Winsorize( - search_space=search_space, - experiment_data=adapter.get_training_data(), - adapter=adapter, - config={"derelativize_with_raw_status_quo": True}, - ) + with self.assertWarnsRegex( + AxOptimizationWarning, "Automatic winsorization isn't supported" + ): + t = Winsorize( + search_space=search_space, + experiment_data=adapter.get_training_data(), + adapter=adapter, + config={"derelativize_with_raw_status_quo": True}, + ) self.assertDictEqual( - t.cutoffs, {"a": (-INF, 2.625), "b": (-INF, 4.25), "c": (-3.25, INF)} + t.cutoffs, {"a": (-INF, 4.25), "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 4c2aee262ce..dc79eb4c161 100644 --- a/ax/adapter/transforms/winsorize.py +++ b/ax/adapter/transforms/winsorize.py @@ -28,7 +28,6 @@ ComparisonOp, ObjectiveThreshold, OutcomeConstraint, - ScalarizedOutcomeConstraint, ) from ax.core.search_space import SearchSpace from ax.exceptions.core import AxOptimizationWarning, UserInputError @@ -313,8 +312,7 @@ def _obtain_cutoffs_from_outcome_constraints( """Get all outcome constraints (non-scalarized) for a given metric.""" # Check for scalarized outcome constraints for the given metric if any( - isinstance(oc, ScalarizedOutcomeConstraint) - and metric_signature in oc.metric_names + len(oc.metric_weights) > 1 and metric_signature in oc.metric_signatures for oc in optimization_config.outcome_constraints ): warnings.warn( @@ -341,8 +339,7 @@ def _get_non_scalarized_outcome_constraints( return [ oc for oc in optimization_config.outcome_constraints - if not isinstance(oc, ScalarizedOutcomeConstraint) - and oc.metric_names[0] == metric_signature + if len(oc.metric_weights) == 1 and oc.metric_signatures[0] == metric_signature ] diff --git a/ax/api/tests/test_client.py b/ax/api/tests/test_client.py index 2f0bb5ddbc4..0602b96b662 100644 --- a/ax/api/tests/test_client.py +++ b/ax/api/tests/test_client.py @@ -16,6 +16,7 @@ import numpy as np import pandas as pd +from ax.adapter.adapter_utils import extract_outcome_constraints from ax.analysis.analysis import NOT_APPLICABLE_STATE_SUBTITLE from ax.analysis.plotly.parallel_coordinates import ParallelCoordinatesPlot from ax.api.client import Client @@ -43,6 +44,7 @@ from ax.core.search_space import SearchSpace from ax.core.trial import Trial from ax.core.trial_status import TrialStatus +from ax.core.types import TBounds from ax.core.utils import compute_metric_availability, MetricAvailability from ax.early_stopping.strategies import PercentileEarlyStoppingStrategy from ax.exceptions.core import UnsupportedError @@ -1956,6 +1958,56 @@ def test_summarize_with_derived_parameters(self) -> None: self.assertIn("param_sum", summary_df.columns) self.assertEqual(len(summary_df), 3) + @mock_botorch_optimize + def test_expression_outcome_constraint_generation(self) -> None: + """Keep a multi-metric constraint intact through Client's transform chain.""" + client = Client() + client.configure_experiment( + parameters=[ + RangeParameterConfig(name="x", parameter_type="float", bounds=(0, 1)) + ] + ) + client.configure_optimization( + objective="obj", outcome_constraints=["a + 2 * b <= 1.0"] + ) + client.configure_generation_strategy( + initialization_budget=3, initialization_random_seed=123 + ) + x_values = [] + for _ in range(3): + for index, parameters in client.get_next_trials(max_trials=1).items(): + x = assert_is_instance(parameters["x"], float) + x_values.append(x) + client.complete_trial( + trial_index=index, raw_data={"obj": x, "a": x, "b": 1 - x} + ) + captured: list[tuple[list[str], TBounds]] = [] + + def capture_constraints( + outcome_constraints: list[OutcomeConstraint], outcomes: list[str] + ) -> TBounds: + result = extract_outcome_constraints( + outcome_constraints=outcome_constraints, outcomes=outcomes + ) + captured.append((outcomes, result)) + return result + + with mock.patch( + "ax.adapter.torch.extract_outcome_constraints", + side_effect=capture_constraints, + ): + self.assertEqual(len(client.get_next_trials(max_trials=1)), 1) + self.assertGreater(len(captured), 0) + outcomes, constraints = captured[-1] + A, b = none_throws(constraints) + metrics = client._experiment.metrics + std = np.std(x_values, ddof=1) + expected = np.zeros((1, len(outcomes))) + expected[0, outcomes.index(metrics["a"].signature)] = std + expected[0, outcomes.index(metrics["b"].signature)] = 2 * std + np.testing.assert_allclose(A, expected) + np.testing.assert_allclose(b, [[np.mean(x_values) - 1]]) + @mock_botorch_optimize def test_scalarized_single_objective_optimization(self) -> None: """Test that a scalarized single objective (e.g., "0.5*m1 + 0.5*m2") diff --git a/ax/api/utils/instantiation/from_string.py b/ax/api/utils/instantiation/from_string.py index 4c9468e3894..56d788aee69 100644 --- a/ax/api/utils/instantiation/from_string.py +++ b/ax/api/utils/instantiation/from_string.py @@ -67,6 +67,13 @@ def optimization_config_from_string( else None ) + for constraint in outcome_constraints or []: + if constraint.relative and len(constraint.metric_weights) > 1: + raise UserInputError( + "Scalarized outcome constraints cannot be relative. " + f"Use an absolute bound instead, found {constraint.expression}." + ) + if objective.is_multi_objective: # A single-metric constraint on an objective metric becomes an # objective threshold only when it bounds the objective from its diff --git a/ax/api/utils/instantiation/tests/test_from_string.py b/ax/api/utils/instantiation/tests/test_from_string.py index 6a2dca07a60..b6d9079a0c3 100644 --- a/ax/api/utils/instantiation/tests/test_from_string.py +++ b/ax/api/utils/instantiation/tests/test_from_string.py @@ -109,6 +109,40 @@ def test_constraint_against_optimization_direction_on_objective(self) -> None: ), ) + def test_relative_scalarized_constraints(self) -> None: + """Reject relative multi-metric expressions before generation.""" + for objective in ["obj", "obj, obj2"]: + for expression in [ + "a + 2 * b >= 0.9 * baseline", + "a + 2 * b <= 0.9 * baseline", + "a - 2 * b >= 0.9 * baseline", + ]: + with ( + self.subTest(objective=objective, expression=expression), + self.assertRaisesRegex( + UserInputError, + "Scalarized outcome constraints cannot be relative", + ), + ): + optimization_config_from_string( + objective_str=objective, outcome_constraint_strs=[expression] + ) + for expression in ["a >= 0.9 * baseline", "a + 0 * b >= 0.9 * baseline"]: + config = optimization_config_from_string( + objective_str=objective, outcome_constraint_strs=[expression] + ) + self.assertTrue(config.outcome_constraints[0].relative) + self.assertEqual( + config.outcome_constraints[0].metric_weights, [("a", 1)] + ) + config = optimization_config_from_string( + objective_str=objective, outcome_constraint_strs=["a + 2 * b <= 1.0"] + ) + self.assertFalse(config.outcome_constraints[0].relative) + self.assertCountEqual( + config.outcome_constraints[0].metric_weights, [("a", 1), ("b", 2)] + ) + def test_objective_constraint_on_single_objective_raises(self) -> None: with self.assertRaisesRegex( UserInputError, "Outcome constraints may not be placed"