Preserve multi-metric outcome constraints through candidate generation

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)
 

ax/adapter/tests/test_utils.py

@@ -16,19 +16,22 @@
     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,
     ScalarizedOutcomeConstraint,
 )
 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(

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(

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 = {}

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])

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