Replace OrderConstraint with ParameterConstraint (#4659)

Summary:
Pull Request resolved: #4659

Any OrderConstraint(a, b) can be rewritten ParameterConstraint("a <= b") which is both more clear and removes polymorphism ahead of storage rework

Reviewed By: saitcakmak

Differential Revision: D88890559

Author: mpolson64

ax/adapter/tests/test_random_adapter.py

@@ -18,11 +18,7 @@
 from ax.core.metric import Metric
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ParameterType, RangeParameter
-from ax.core.parameter_constraint import (
-    OrderConstraint,
-    ParameterConstraint,
-    SumConstraint,
-)
+from ax.core.parameter_constraint import ParameterConstraint, SumConstraint
 from ax.core.search_space import SearchSpace
 from ax.exceptions.core import SearchSpaceExhausted
 from ax.generators.random.base import RandomGenerator
@@ -43,7 +39,7 @@ def setUp(self) -> None:
         z = RangeParameter("z", ParameterType.FLOAT, lower=0, upper=5)
         self.parameters = [x, y, z]
         parameter_constraints: list[ParameterConstraint] = [
-            OrderConstraint(x, y),
+            ParameterConstraint(inequality="x <= y"),
             SumConstraint([x, z], False, 3.5),
         ]
         self.search_space = SearchSpace(self.parameters, parameter_constraints)

ax/adapter/transforms/rounding.py

@@ -12,7 +12,6 @@
 
 import numpy as np
 import numpy.typing as npt
-from ax.core.parameter_constraint import OrderConstraint
 from ax.core.search_space import SearchSpace
 from ax.core.types import TParameterization
 
@@ -56,13 +55,8 @@ def strict_onehot_round(x: npt.NDArray) -> npt.NDArray:
 def contains_constrained_integer(
     search_space: SearchSpace, transform_parameters: set[str]
 ) -> bool:
-    """Check if any integer parameters are present in parameter_constraints.
-
-    Order constraints are ignored since strict rounding preserves ordering.
-    """
+    """Check if any integer parameters are present in parameter_constraints."""
     for constraint in search_space.parameter_constraints:
-        if isinstance(constraint, OrderConstraint):
-            continue
         constraint_params = set(constraint.constraint_dict.keys())
         if constraint_params.intersection(transform_parameters):
             return True

ax/adapter/transforms/tests/test_fixed_to_tunable.py

@@ -18,7 +18,7 @@
     ParameterType,
     RangeParameter,
 )
-from ax.core.parameter_constraint import OrderConstraint
+from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_experiment_with_observations
@@ -120,11 +120,7 @@ def test_transform_search_space_with_constraints(self) -> None:
         ]
         search_space_with_constraints = SearchSpace(
             parameters=parameters,
-            parameter_constraints=[
-                OrderConstraint(
-                    lower_parameter=parameters[0], upper_parameter=parameters[1]
-                )
-            ],
+            parameter_constraints=[ParameterConstraint(inequality="x <= y")],
         )
 
         # Joint space with range parameter for 'y'

ax/adapter/transforms/tests/test_int_to_float_transform.py

@@ -15,7 +15,7 @@
 from ax.adapter.transforms.int_to_float import IntToFloat
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
-from ax.core.parameter_constraint import OrderConstraint, SumConstraint
+from ax.core.parameter_constraint import ParameterConstraint, SumConstraint
 from ax.core.search_space import SearchSpace
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_experiment_with_observations
@@ -36,11 +36,7 @@ def setUp(self) -> None:
         ]
         self.search_space = SearchSpace(
             parameters=parameters,
-            parameter_constraints=[
-                OrderConstraint(
-                    lower_parameter=parameters[0], upper_parameter=parameters[1]
-                )
-            ],
+            parameter_constraints=[ParameterConstraint(inequality="x <= a")],
         )
         self.t = IntToFloat(search_space=self.search_space)
         self.t2 = IntToFloat(

ax/core/__init__.py

@@ -31,11 +31,7 @@
     ParameterType,
     RangeParameter,
 )
-from ax.core.parameter_constraint import (
-    OrderConstraint,
-    ParameterConstraint,
-    SumConstraint,
-)
+from ax.core.parameter_constraint import ParameterConstraint, SumConstraint
 from ax.core.runner import Runner
 from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
@@ -57,7 +53,6 @@
     "ObjectiveThreshold",
     "ObservationFeatures",
     "OptimizationConfig",
-    "OrderConstraint",
     "OutcomeConstraint",
     "Parameter",
     "ParameterConstraint",

ax/core/parameter_constraint.py

@@ -114,70 +114,6 @@ def _unique_id(self) -> str:
         return str(self)
 
 
-class OrderConstraint(ParameterConstraint):
-    """Constraint object for specifying one parameter to be smaller than another."""
-
-    _bound: float
-
-    def __init__(self, lower_parameter: Parameter, upper_parameter: Parameter) -> None:
-        """Initialize OrderConstraint
-
-        Args:
-            lower_parameter: Parameter that should have the lower value.
-            upper_parameter: Parameter that should have the higher value.
-
-        Note:
-            The constraint p1 <= p2 can be expressed in matrix notation as
-            [1, -1] * [p1, p2]^T <= 0.
-        """
-        validate_constraint_parameters([lower_parameter, upper_parameter])
-
-        self._lower_parameter = lower_parameter
-        self._upper_parameter = upper_parameter
-        self._bound = 0.0
-
-    @property
-    def lower_parameter(self) -> Parameter:
-        """Parameter with lower value."""
-        return self._lower_parameter
-
-    @property
-    def upper_parameter(self) -> Parameter:
-        """Parameter with higher value."""
-        return self._upper_parameter
-
-    @property
-    def parameters(self) -> list[Parameter]:
-        """Parameters."""
-        return [self.lower_parameter, self.upper_parameter]
-
-    @property
-    def constraint_dict(self) -> dict[str, float]:
-        """Weights on parameters for linear constraint representation."""
-        return {self.lower_parameter.name: 1.0, self.upper_parameter.name: -1.0}
-
-    def clone(self) -> OrderConstraint:
-        """Clone."""
-        return OrderConstraint(
-            lower_parameter=self.lower_parameter.clone(),
-            upper_parameter=self._upper_parameter.clone(),
-        )
-
-    def clone_with_transformed_parameters(
-        self, transformed_parameters: dict[str, Parameter]
-    ) -> OrderConstraint:
-        """Clone, but replace parameters with transformed versions."""
-        return OrderConstraint(
-            lower_parameter=transformed_parameters[self.lower_parameter.name],
-            upper_parameter=transformed_parameters[self._upper_parameter.name],
-        )
-
-    def __repr__(self) -> str:
-        return "OrderConstraint({} <= {})".format(
-            self.lower_parameter.name, self.upper_parameter.name
-        )
-
-
 class SumConstraint(ParameterConstraint):
     """Constraint on the sum of parameters being greater or less than a bound."""
 
@@ -281,6 +217,7 @@ def validate_constraint_parameters(parameters: Sequence[Parameter]) -> None:
         if not isinstance(parameter, RangeParameter):
             raise ValueError(
                 "All parameters in a parameter constraint must be RangeParameters."
+                f"Found {parameter}"
             )
 
         # Log parameters require a non-linear transformation, and Ax

ax/core/search_space.py

@@ -27,9 +27,9 @@
     TParamValue,
 )
 from ax.core.parameter_constraint import (
-    OrderConstraint,
     ParameterConstraint,
     SumConstraint,
+    validate_constraint_parameters,
 )
 from ax.core.types import TParameterization
 from ax.exceptions.core import AxWarning, UnsupportedError, UserInputError
@@ -194,17 +194,16 @@ def set_parameter_constraints(
         # the matching name among the search space's parameters, so we
         # are not keeping two copies of the same parameter.
         for constraint in parameter_constraints:
-            if isinstance(constraint, OrderConstraint):
-                constraint._lower_parameter = self.parameters[
-                    constraint._lower_parameter.name
-                ]
-                constraint._upper_parameter = self.parameters[
-                    constraint._upper_parameter.name
-                ]
-            elif isinstance(constraint, SumConstraint):
+            if isinstance(constraint, SumConstraint):
                 for idx, parameter in enumerate(constraint.parameters):
                     constraint.parameters[idx] = self.parameters[parameter.name]
 
+            validate_constraint_parameters(
+                parameters=[
+                    self._parameters[name] for name in constraint.constraint_dict.keys()
+                ]
+            )
+
         self._parameter_constraints: list[ParameterConstraint] = parameter_constraints
 
     def add_parameters(
@@ -534,9 +533,7 @@ def _validate_parameter_constraints(
         self, parameter_constraints: list[ParameterConstraint]
     ) -> None:
         for constraint in parameter_constraints:
-            if isinstance(constraint, OrderConstraint) or isinstance(
-                constraint, SumConstraint
-            ):
+            if isinstance(constraint, SumConstraint):
                 for parameter in constraint.parameters:
                     if parameter.name not in self._parameters.keys():
                         raise ValueError(

ax/core/tests/test_parameter_constraint.py

@@ -6,15 +6,9 @@
 
 # pyre-strict
 
-from ax.core.parameter import (
-    ChoiceParameter,
-    FixedParameter,
-    ParameterType,
-    RangeParameter,
-)
+from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import (
     ComparisonOp,
-    OrderConstraint,
     ParameterConstraint,
     SumConstraint,
 )
@@ -135,62 +129,6 @@ def test_Sortable(self) -> None:
         self.assertTrue(constraint1 < constraint2)
 
 
-class OrderConstraintTest(TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-        self.x = RangeParameter("x", ParameterType.INT, lower=0, upper=1)
-        self.y = RangeParameter("y", ParameterType.INT, lower=0, upper=1)
-        self.constraint = OrderConstraint(
-            lower_parameter=self.x, upper_parameter=self.y
-        )
-        self.constraint_repr = "OrderConstraint(x <= y)"
-
-    def test_Properties(self) -> None:
-        self.assertEqual(self.constraint.lower_parameter.name, "x")
-        self.assertEqual(self.constraint.upper_parameter.name, "y")
-
-    def test_Repr(self) -> None:
-        self.assertEqual(str(self.constraint), self.constraint_repr)
-
-    def test_Validate(self) -> None:
-        self.assertTrue(self.constraint.check({"x": 0, "y": 1}))
-        self.assertTrue(self.constraint.check({"x": 1, "y": 1}))
-        self.assertFalse(self.constraint.check({"x": 1, "y": 0}))
-
-    def test_Clone(self) -> None:
-        constraint_clone = self.constraint.clone()
-        self.assertEqual(
-            self.constraint.lower_parameter, constraint_clone.lower_parameter
-        )
-
-        constraint_clone._lower_parameter = self.y
-        self.assertNotEqual(
-            self.constraint.lower_parameter, constraint_clone.lower_parameter
-        )
-
-    def test_CloneWithTransformedParameters(self) -> None:
-        constraint_clone = self.constraint.clone_with_transformed_parameters(
-            transformed_parameters={p.name: p for p in self.constraint.parameters}
-        )
-        self.assertEqual(
-            self.constraint.lower_parameter, constraint_clone.lower_parameter
-        )
-
-        constraint_clone._lower_parameter = self.y
-        self.assertNotEqual(
-            self.constraint.lower_parameter, constraint_clone.lower_parameter
-        )
-
-    def test_InvalidSetup(self) -> None:
-        z = FixedParameter("z", ParameterType.INT, 0)
-        with self.assertRaises(ValueError):
-            self.constraint = OrderConstraint(lower_parameter=self.x, upper_parameter=z)
-
-        z = ChoiceParameter("z", ParameterType.STRING, ["a", "b", "c"])
-        with self.assertRaises(ValueError):
-            self.constraint = OrderConstraint(lower_parameter=self.x, upper_parameter=z)
-
-
 class SumConstraintTest(TestCase):
     def setUp(self) -> None:
         super().setUp()