Move metric_name_to_signature from Adapter to Objective/OutcomeConstraint (facebook#5122)

Summary:
Pull Request resolved: facebook#5122

The `metric_name_to_signature` mapping (user-facing metric name -> canonical signature used in data/modeling) previously lived on `Adapter`, built from `experiment.metrics` at init time. This was architecturally awkward because the mapping is consumed alongside `OptimizationConfig` objects (objectives, constraints) throughout adapter_utils, transforms, and cross-validation, but lived on a separate object.

This diff co-locates the mapping with the objects that reference metrics by name:

**Core changes (ax/core):**

- `Objective` and `OutcomeConstraint` now accept an optional `metric_name_to_signature` kwarg, stored as an instance field
- Both classes expose `metric_signatures` and `metric_name_to_signature` properties (with getter and setter)
- `metric_weights` on both classes now returns `(signature, weight)` tuples instead of `(name, weight)`, aligning with how downstream consumers use them
- `OptimizationConfig` aggregates these into `metric_signatures` and `metric_name_to_signature` properties, with a setter that propagates to children
- `clone()` methods preserve the mapping

**Adapter changes (ax/adapter):**

- Removed `Adapter._metric_name_to_signature` field and property
- Removed `metric_name_to_signature` parameter from `adapter_utils.py` functions: `extract_objective_thresholds`, `extract_objective_weights`, `extract_objective_weight_matrix`, `extract_outcome_constraints`, `feasible_hypervolume`
- Removed `metric_name_to_signature` parameter from `has_good_opt_config_model_fit` in `cross_validation.py`
- Removed `metric_name_to_signature` parameter from `validate_transformed_optimization_config` in `torch.py`
- Updated all transforms to use `constraint.metric_signatures` and `objective.metric_signatures` directly instead of looking up through `adapter.metric_name_to_signature`

**Orchestration test fix (ax/orchestration):**

- Updated `test_get_best_trial` in `test_orchestrator.py` to expect `ValueError` instead of `KeyError` when passing a mismatched `OptimizationConfig`. Previously the dict lookup on `metric_name_to_signature` raised `KeyError`; now `outcomes.index()` raises `ValueError`.

**Storage:** No changes needed -- the mapping is not persisted on its own and instead is reconstructed by Experiment initialization.

Reviewed By: saitcakmak

Differential Revision: D98837790

fbshipit-source-id: b4dbef567a1601f7f23909fac8f8f6383cd11b09

Author: mpolson64

ax/adapter/adapter_utils.py

@@ -210,7 +210,6 @@ def extract_objective_thresholds(
     objective_thresholds: TRefPoint,
     objective: Objective,
     outcomes: list[str],
-    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray | None:
     """Extracts objective thresholds' values, in the order of objectives.
 
@@ -231,8 +230,7 @@ def extract_objective_thresholds(
     Args:
         objective_thresholds: Objective thresholds to extract values from.
         objective: The corresponding Objective, for validation purposes.
-        outcomes: n-length list of names of metrics.
-        metric_name_to_signature: Mapping from metric names to signatures.
+        outcomes: n-length list of metric signatures.
 
     Returns:
         ``(n_objectives,)`` array of maximization-aligned thresholds, or None.
@@ -242,7 +240,7 @@ def extract_objective_thresholds(
 
     objective_threshold_dict = {}
     for ot in objective_thresholds:
-        ot_signature = metric_name_to_signature[ot.metric_names[0]]
+        ot_signature = ot.metric_signatures[0]
         if ot.relative:
             raise ValueError(
                 f"Objective {ot_signature} has a relative threshold that "
@@ -251,9 +249,7 @@ def extract_objective_thresholds(
         objective_threshold_dict[ot_signature] = ot.bound
 
     # Check that all thresholds correspond to a metric.
-    obj_metric_signatures = [
-        metric_name_to_signature[name] for name in objective.metric_names
-    ]
+    obj_metric_signatures = objective.metric_signatures
     if set(objective_threshold_dict.keys()).difference(set(obj_metric_signatures)):
         raise ValueError(
             "Some objective thresholds do not have corresponding metrics. "
@@ -273,7 +269,7 @@ def extract_objective_thresholds(
         if len(sub_mw) > 1:
             continue  # Scalarized sub-objective — NaN, will be inferred later.
         name, weight = sub_mw[0]
-        sig = metric_name_to_signature[name]
+        sig = objective.metric_name_to_signature[name]
         if sig in objective_threshold_dict:
             sign = 1.0 if weight > 0 else -1.0
             obj_t[i] = sign * objective_threshold_dict[sig]
@@ -283,7 +279,6 @@ def extract_objective_thresholds(
 def extract_objective_weights(
     objective: Objective,
     outcomes: list[str],
-    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray:
     """Extract a weights for objectives.
 
@@ -301,25 +296,22 @@ def extract_objective_weights(
     Args:
         objective: Objective to extract weights from.
         outcomes: n-length list of metric signatures.
-        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns:
         n-length array of weights.
 
     """
     objective_weights = np.zeros(len(outcomes))
-    # metric_weights returns sign-encoded (name, weight) tuples for all
+    # metric_weights returns sign-encoded (signature, weight) tuples for all
     # objective types (single, scalarized, multi).
-    for obj_metric_name, obj_weight in objective.metric_weights:
-        sig = metric_name_to_signature[obj_metric_name]
-        objective_weights[outcomes.index(sig)] = obj_weight
+    for obj_metric_sig, obj_weight in objective.metric_weights:
+        objective_weights[outcomes.index(obj_metric_sig)] = obj_weight
     return objective_weights
 
 
 def extract_objective_weight_matrix(
     objective: Objective,
     outcomes: list[str],
-    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray:
     """Extract a 2D weight matrix for objectives.
 
@@ -333,20 +325,24 @@ def extract_objective_weight_matrix(
 
     Args:
         objective: Objective to extract weights from.
-        outcomes: n-length list of signatures of metrics.
-        metric_name_to_signature: Mapping from metric names to signatures.
+        outcomes: n-length list of metric signatures.
 
     Returns:
         ``(n_objectives, n)`` array of weights.
     """
     if objective.is_multi_objective:
         rows: list[npt.NDArray] = []
-        for name, weight in objective.metric_weights:
+        obj_names = objective.metric_names
+        obj_weights = [w for _, w in objective.metric_weights]
+        name_to_sig = objective.metric_name_to_signature
+        for name, weight in zip(obj_names, obj_weights):
             rows.append(
                 extract_objective_weights(
-                    objective=Objective(expression=f"{weight} * {name}"),
+                    objective=Objective(
+                        expression=f"{weight} * {name}",
+                        metric_name_to_signature={name: name_to_sig[name]},
+                    ),
                     outcomes=outcomes,
-                    metric_name_to_signature=metric_name_to_signature,
                 )
             )
         return np.stack(rows, axis=0)
@@ -355,14 +351,12 @@ def extract_objective_weight_matrix(
         return extract_objective_weights(
             objective=objective,
             outcomes=outcomes,
-            metric_name_to_signature=metric_name_to_signature,
         ).reshape(1, -1)
 
 
 def extract_outcome_constraints(
     outcome_constraints: list[OutcomeConstraint],
     outcomes: list[str],
-    metric_name_to_signature: Mapping[str, str],
 ) -> TBounds:
     if len(outcome_constraints) == 0:
         return None
@@ -372,11 +366,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_name, c_weight in c.metric_weights:
-                j = outcomes.index(metric_name_to_signature[c_metric_name])
+            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(metric_name_to_signature[c.metric_names[0]])
+            j = outcomes.index(c.metric_signatures[0])
             A[i, j] = s
         b[i, 0] = s * c.bound
     return (A, b)
@@ -689,18 +683,15 @@ def get_pareto_frontier_and_configs(
     objective_weights = extract_objective_weight_matrix(
         objective=optimization_config.objective,
         outcomes=adapter.outcomes,
-        metric_name_to_signature=adapter.metric_name_to_signature,
     )
     outcome_constraints = extract_outcome_constraints(
         outcome_constraints=optimization_config.outcome_constraints,
         outcomes=adapter.outcomes,
-        metric_name_to_signature=adapter.metric_name_to_signature,
     )
     obj_t = extract_objective_thresholds(
         objective_thresholds=optimization_config.objective_thresholds,
         objective=optimization_config.objective,
         outcomes=adapter.outcomes,
-        metric_name_to_signature=adapter.metric_name_to_signature,
     )
     if obj_t is not None:
         obj_t = array_to_tensor(obj_t)
@@ -1155,26 +1146,23 @@ def observation_features_to_array(
 def feasible_hypervolume(
     optimization_config: MultiObjectiveOptimizationConfig,
     values: dict[str, npt.NDArray],
-    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray:
     """Compute the feasible hypervolume each iteration.
 
     Args:
         optimization_config: Optimization config.
-        values: Dictionary from metric name to array of value at each
+        values: Dictionary from metric signature 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`.
-        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns: Array of feasible hypervolumes.
     """
     # Get objective at each iteration
     obj_threshold_dict = {
-        metric_name_to_signature[ot.metric_names[0]]: ot.bound
+        ot.metric_signatures[0]: ot.bound
         for ot in optimization_config.objective_thresholds
     }
-    obj_metric_names = optimization_config.objective.metric_names
-    obj_metric_sigs = [metric_name_to_signature[name] for name in obj_metric_names]
+    obj_metric_sigs = optimization_config.objective.metric_signatures
     f_vals = np.hstack([values[sig].reshape(-1, 1) for sig in obj_metric_sigs])
     obj_thresholds = np.array([obj_threshold_dict[sig] for sig in obj_metric_sigs])
     # Set infeasible points to be the objective threshold
@@ -1183,7 +1171,7 @@ def feasible_hypervolume(
             raise ValueError(
                 "Benchmark aggregation does not support relative constraints"
             )
-        oc_sig = metric_name_to_signature[oc.metric_names[0]]
+        oc_sig = oc.metric_signatures[0]
         g = values[oc_sig]
         feas = g <= oc.bound if oc.op == ComparisonOp.LEQ else g >= oc.bound
         f_vals[~feas] = obj_thresholds
@@ -1192,7 +1180,7 @@ def feasible_hypervolume(
     # Positive weight = maximize, negative weight = minimize.
     obj_weight_dict = dict(optimization_config.objective.metric_weights)
     obj_weights = np.array(
-        [1 if obj_weight_dict[name] > 0 else -1 for name in obj_metric_names]
+        [1 if obj_weight_dict[sig] > 0 else -1 for sig in obj_metric_sigs]
     )
     obj_thresholds = obj_thresholds * obj_weights
     f_vals = f_vals * obj_weights

ax/adapter/base.py

@@ -191,9 +191,6 @@ def __init__(
         )
         self._experiment_properties: dict[str, Any] = experiment._properties
         self._experiment: Experiment = experiment
-        self._metric_name_to_signature: dict[str, str] = {
-            name: metric.signature for name, metric in self._experiment.metrics.items()
-        }
 
         if self._optimization_config is None:
             self._optimization_config = experiment.optimization_config
@@ -563,11 +560,6 @@ def metric_signatures(self) -> set[str]:
         """Metric signatures present in training data."""
         return self._metric_signatures
 
-    @property
-    def metric_name_to_signature(self) -> dict[str, str]:
-        """Mapping from metric names to their signatures."""
-        return self._metric_name_to_signature
-
     @property
     def model_space(self) -> SearchSpace:
         """SearchSpace used to fit model."""

ax/adapter/cross_validation.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from collections.abc import Callable, Iterable, Mapping
+from collections.abc import Callable, Iterable
 from logging import Logger
 from typing import cast, NamedTuple
 from warnings import warn
@@ -573,7 +573,6 @@ def assess_model_fit(
 def has_good_opt_config_model_fit(
     optimization_config: OptimizationConfig,
     assess_model_fit_result: AssessModelFitResult,
-    metric_name_to_signature: Mapping[str, str],
 ) -> bool:
     """Assess model fit for given diagnostics results across the optimization
     config metrics
@@ -586,7 +585,6 @@ def has_good_opt_config_model_fit(
     Args:
         optimization_config: Objective/Outcome constraint metrics to assess
         assess_model_fit_result: Output of assess_model_fit
-        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns:
         Two dictionaries, one for good metrics, one for bad metrics, each
@@ -596,9 +594,8 @@ 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(
-        metric_name_to_signature[name]
-        in assess_model_fit_result.good_fit_metrics_to_fisher_score
-        for name in optimization_config.objective.metric_names
+        sig in assess_model_fit_result.good_fit_metrics_to_fisher_score
+        for sig in optimization_config.objective.metric_signatures
     )
     return has_good_opt_config_fit
 

ax/adapter/discrete.py

@@ -166,17 +166,14 @@ def _gen(
             validate_transformed_optimization_config(
                 optimization_config,
                 self.outcomes,
-                metric_name_to_signature=self.metric_name_to_signature,
             )
             objective_weights = extract_objective_weights(
                 objective=optimization_config.objective,
                 outcomes=self.outcomes,
-                metric_name_to_signature=self.metric_name_to_signature,
             )
             outcome_constraints = extract_outcome_constraints(
                 outcome_constraints=optimization_config.outcome_constraints,
                 outcomes=self.outcomes,
-                metric_name_to_signature=self.metric_name_to_signature,
             )
 
         # Get fixed features

ax/adapter/tests/test_adapter_utils.py

@@ -78,8 +78,6 @@ def test_feasible_hypervolume(self) -> None:
                 ),
             ],
         )
-        # For plain Metric objects, signature == name.
-        metric_name_to_signature = {m.name: m.name for m in [ma, mb, mc]}
         feas_hv = feasible_hypervolume(
             optimization_config,
             values={
@@ -108,7 +106,6 @@ def test_feasible_hypervolume(self) -> None:
                     ]
                 ),
             },
-            metric_name_to_signature=metric_name_to_signature,
         )
         self.assertEqual(list(feas_hv), [0.0, 0.0, 1.0, 1.0])
 
@@ -540,15 +537,12 @@ 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"]
-        # For plain Metric objects, signature == name.
-        metric_name_to_signature = {name: name for name in outcomes}
 
         # Single Objective: one row, nonzero only in matching column.
         obj = Objective(metric=m1, minimize=False)
         result = extract_objective_weight_matrix(
             objective=obj,
             outcomes=outcomes,
-            metric_name_to_signature=metric_name_to_signature,
         )
         np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0]])
 
@@ -557,7 +551,6 @@ def test_extract_objective_weight_matrix(self) -> None:
         result = extract_objective_weight_matrix(
             objective=obj_min,
             outcomes=outcomes,
-            metric_name_to_signature=metric_name_to_signature,
         )
         np.testing.assert_array_equal(result, [[0.0, -1.0, 0.0]])
 
@@ -566,7 +559,6 @@ def test_extract_objective_weight_matrix(self) -> None:
         result = extract_objective_weight_matrix(
             objective=scal,
             outcomes=outcomes,
-            metric_name_to_signature=metric_name_to_signature,
         )
         np.testing.assert_array_almost_equal(result, [[0.3, 0.0, 0.7]])
 
@@ -580,7 +572,6 @@ def test_extract_objective_weight_matrix(self) -> None:
         result = extract_objective_weight_matrix(
             objective=multi,
             outcomes=outcomes,
-            metric_name_to_signature=metric_name_to_signature,
         )
         np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0]])
 

ax/adapter/tests/test_cross_validation.py

@@ -480,7 +480,6 @@ 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,
-            metric_name_to_signature={"m1": "m1", "m2": "m2"},
         )
         self.assertFalse(has_good_fit)
 
@@ -496,7 +495,6 @@ 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,
-            metric_name_to_signature={"m1": "m1", "m2": "m2"},
         )
         self.assertFalse(has_good_fit)
 
@@ -510,7 +508,6 @@ 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,
-            metric_name_to_signature={"m1": "m1", "m2": "m2"},
         )
         self.assertFalse(has_good_fit)
 

ax/adapter/tests/test_torch_moo_adapter.py

@@ -464,7 +464,10 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
         first = sub_exprs[0]
         if not first.startswith("-"):
             sub_exprs[0] = f"-{first}"
-        oc.objective = Objective(expression=", ".join(sub_exprs))
+        oc.objective = Objective(
+            expression=", ".join(sub_exprs),
+            metric_name_to_signature={s.lstrip("-"): s.lstrip("-") for s in sub_exprs},
+        )
 
         for use_partial_thresholds in (False, True):
             if use_partial_thresholds: