Make ParameterConstraint store its own inequality string (#4658)

Summary:

Move some of the string parsing from ax/api into core Ax to allow a ParameterConstraint to be defined and constructed from its inequality string.

This will allow more flexibility in the future, but is primarily motivated by a want to store only the inequality string when saving ParameterConstraints and to eliminiate unnnecessary subclasses SumConstraint and OrderConstraint

Differential Revision: D88880607

Author: mpolson64

ax/adapter/tests/test_torch_moo_adapter.py

@@ -437,9 +437,7 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
         )
         fixed_features = ObservationFeatures(parameters={"x1": 0.0})
         search_space = exp.search_space.clone()
-        param_constraints = [
-            ParameterConstraint(constraint_dict={"x1": 1.0}, bound=10.0)
-        ]
+        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

ax/adapter/transforms/tests/test_choice_encode_transform.py

@@ -68,9 +68,7 @@ def setUp(self) -> None:
                     sort_values=False,
                 ),
             ],
-            parameter_constraints=[
-                ParameterConstraint(constraint_dict={"x": -0.5, "a": 1}, bound=0.5)
-            ],
+            parameter_constraints=[ParameterConstraint(inequality="-0.5*x + a <= 0.5")],
         )
         self.t = self.t_class(search_space=self.search_space)
         input_params: TParameterization = {

ax/adapter/transforms/tests/test_one_hot_transform.py

@@ -46,9 +46,7 @@ def setUp(self) -> None:
                     is_ordered=True,
                 ),
             ],
-            parameter_constraints=[
-                ParameterConstraint(constraint_dict={"x": -0.5, "a": 1}, bound=0.5)
-            ],
+            parameter_constraints=[ParameterConstraint(inequality="-0.5*x + a <= 0.5")],
         )
         self.t = OneHot(search_space=self.search_space)
         self.t2 = OneHot(

ax/adapter/transforms/tests/test_unit_x_transform.py

@@ -46,8 +46,8 @@ def setUp(self) -> None:
                 ),
             ],
             parameter_constraints=[
-                ParameterConstraint(constraint_dict={"x": -0.5, "y": 1}, bound=0.5),
-                ParameterConstraint(constraint_dict={"x": -0.5, "a": 1}, bound=0.5),
+                ParameterConstraint(inequality="-0.5*x + y <= 0.5"),
+                ParameterConstraint(inequality="-0.5*x + a <= 0.5"),
             ],
         )
         self.t = UnitX(search_space=self.search_space)
@@ -157,8 +157,8 @@ def test_TransformNewSearchSpace(self) -> None:
                 ),
             ],
             parameter_constraints=[
-                ParameterConstraint(constraint_dict={"x": -0.5, "y": 1}, bound=0.5),
-                ParameterConstraint(constraint_dict={"x": -0.5, "a": 1}, bound=0.5),
+                ParameterConstraint(inequality="-0.5*x + y <= 0.5"),
+                ParameterConstraint(inequality="-0.5*x + a <= 0.5"),
             ],
         )
         self.t.transform_search_space(new_ss)

ax/adapter/transforms/unit_x.py

@@ -98,8 +98,14 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                     bound -= w * l
                 else:
                     constraint_dict[p_name] = w
+
+            expr = " + ".join(
+                f"{coeff} * {param}" for param, coeff in constraint_dict.items()
+            )
             new_constraints.append(
-                ParameterConstraint(constraint_dict=constraint_dict, bound=bound)
+                ParameterConstraint(
+                    inequality=f"{expr} <= {bound}",
+                )
             )
         search_space.set_parameter_constraints(new_constraints)
         return search_space

ax/analysis/healthcheck/tests/test_search_space_analysis.py

@@ -97,10 +97,7 @@ def test_search_space_boundary_proportions(self) -> None:
                 ),
             ],
             parameter_constraints=[
-                ParameterConstraint(
-                    constraint_dict={"float_range_1": 1.0, "float_range_2": 1.0},
-                    bound=4.0,
-                )
+                ParameterConstraint(inequality="float_range_1 + float_range_2 <= 4")
             ],
         )
 

ax/api/tests/test_client.py

@@ -111,7 +111,7 @@ def test_configure_experiment(self) -> None:
                     ],
                     parameter_constraints=[
                         ParameterConstraint(
-                            constraint_dict={"int_param": 1, "float_param": -1}, bound=0
+                            inequality="int_param <= float_param",
                         )
                     ],
                 ),

ax/api/utils/instantiation/from_string.py

@@ -8,7 +8,6 @@
 from collections.abc import Sequence
 
 from ax.core.map_metric import MapMetric
-
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
@@ -20,14 +19,13 @@
     OutcomeConstraint,
     ScalarizedOutcomeConstraint,
 )
-from ax.core.parameter_constraint import ParameterConstraint
 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.relational import GreaterThan, LessThan
 from sympy.core.symbol import Symbol
 from sympy.core.sympify import sympify
 
@@ -95,35 +93,6 @@ def optimization_config_from_string(
     )
 
 
-def parse_parameter_constraint(constraint_str: str) -> ParameterConstraint:
-    """
-    Parse a parameter constraint string into a ParameterConstraint object using SymPy.
-    Currently only supports linear constraints of the form "a * x + b * y >= k" or
-    "a * x + b * y <= k".
-    """
-    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 = {}
-    bound = 0
-    for term, coefficient in coefficient_dict.items():
-        if term.is_symbol:
-            constraint_dict[unsanitize_name(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(
-                "Only linear inequality parameter constraints are supported, found "
-                f"{constraint_str}"
-            )
-
-    return ParameterConstraint(constraint_dict=constraint_dict, bound=bound)
-
-
 def parse_objective(objective_str: str) -> Objective:
     """
     Parse an objective string into an Objective object using SymPy.
@@ -154,7 +123,7 @@ def parse_outcome_constraint(constraint_str: str) -> OutcomeConstraint:
     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(
+    coefficient_dict = extract_coefficient_dict_from_inequality(
         inequality_str=constraint_str
     )
 
@@ -248,31 +217,3 @@ def _create_single_objective(expression: Expr) -> Objective:
         )
 
     raise UserInputError(f"Only linear objectives are supported, found {expression}.")
-
-
-def _extract_coefficient_dict_from_inequality(
-    inequality_str: str,
-) -> dict[Symbol, float]:
-    """
-    Use SymPy to parse a string into an inequality, invert if necessary to enforce a
-    less-than relationship, move all terms to the left side, and return the
-    coefficients as a dictionary. This is useful for parsing parameter and outcome
-    constraints.
-    """
-    # Parse the constraint string into a SymPy inequality
-    inequality = sympify(sanitize_name(inequality_str))
-
-    # Check the SymPy object is a valid inequality
-    if not isinstance(inequality, GreaterThan | LessThan):
-        raise UserInputError(f"Expected an inequality, found {inequality_str}")
-
-    # Move all terms to the left side of the inequality and invert if necessary to
-    # enforce a less-than relationship
-    if isinstance(inequality, LessThan):
-        expression = inequality.lhs - inequality.rhs
-    else:
-        expression = inequality.rhs - inequality.lhs
-
-    return {
-        key: float(value) for key, value in expression.as_coefficients_dict().items()
-    }

ax/api/utils/instantiation/from_struct.py

@@ -6,11 +6,13 @@
 # pyre-strict
 
 from ax.api.utils.instantiation.from_config import parameter_from_config
-from ax.api.utils.instantiation.from_string import parse_parameter_constraint
 from ax.api.utils.structs import ExperimentStruct
 from ax.core.evaluations_to_data import DataType
 from ax.core.experiment import Experiment
-from ax.core.parameter_constraint import validate_constraint_parameters
+from ax.core.parameter_constraint import (
+    ParameterConstraint,
+    validate_constraint_parameters,
+)
 from ax.core.search_space import SearchSpace
 
 
@@ -22,8 +24,8 @@ def experiment_from_struct(struct: ExperimentStruct) -> Experiment:
     ]
 
     constraints = [
-        parse_parameter_constraint(constraint_str=constraint_str)
-        for constraint_str in struct.parameter_constraints
+        ParameterConstraint(inequality=inequality)
+        for inequality in struct.parameter_constraints
     ]
 
     # Ensure that all ParameterConstraints are valid and acting on existing parameters

ax/api/utils/instantiation/tests/test_from_config.py

@@ -278,7 +278,7 @@ def test_experiment_from_config(self) -> None:
                     ],
                     parameter_constraints=[
                         ParameterConstraint(
-                            constraint_dict={"int_param": 1, "float_param": -1}, bound=0
+                            inequality="int_param <= float_param",
                         )
                     ],
                 ),
@@ -340,7 +340,7 @@ def test_experiment_from_config(self) -> None:
                     ],
                     parameter_constraints=[
                         ParameterConstraint(
-                            constraint_dict={"int_param": 1, "float_param": -1}, bound=0
+                            inequality="int_param <= float_param",
                         )
                     ],
                 ),