Replace SumConstraint with ParameterConstraint (#4660)

Summary:
Pull Request resolved: #4660

Same as D88890559 but for SumConstraint

Reviewed By: saitcakmak

Differential Revision: D88897672

Author: mpolson64

ax/adapter/tests/test_base_adapter.py

@@ -41,7 +41,7 @@
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.outcome_constraint import ComparisonOp, OutcomeConstraint
 from ax.core.parameter import ParameterType, RangeParameter
-from ax.core.parameter_constraint import SumConstraint
+from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace
 from ax.core.types import TParameterization
 from ax.core.utils import get_target_trial_index
@@ -900,15 +900,7 @@ def test_set_model_space(self) -> None:
         trial.mark_completed()
         # Make search space with a parameter constraint
         ss = experiment.search_space.clone()
-        ss.set_parameter_constraints(
-            [
-                SumConstraint(
-                    parameters=list(ss.parameters.values()),
-                    is_upper_bound=True,
-                    bound=30.0,
-                )
-            ]
-        )
+        ss.set_parameter_constraints([ParameterConstraint(inequality="x1 + x2 <= 30")])
 
         # Check that SQ and custom are OOD
         m = Adapter(

ax/adapter/tests/test_random_adapter.py

@@ -18,7 +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 ParameterConstraint, SumConstraint
+from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace
 from ax.exceptions.core import SearchSpaceExhausted
 from ax.generators.random.base import RandomGenerator
@@ -40,7 +40,7 @@ def setUp(self) -> None:
         self.parameters = [x, y, z]
         parameter_constraints: list[ParameterConstraint] = [
             ParameterConstraint(inequality="x <= y"),
-            SumConstraint([x, z], False, 3.5),
+            ParameterConstraint(inequality="x + z >= 3.5"),
         ]
         self.search_space = SearchSpace(self.parameters, parameter_constraints)
         self.experiment = Experiment(search_space=self.search_space)

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 ParameterConstraint, SumConstraint
+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
@@ -239,11 +239,7 @@ def test_RoundingWithConstrainedIntRanges(self) -> None:
         ]
         constrained_int_search_space = SearchSpace(
             parameters=parameters,
-            parameter_constraints=[
-                # pyre-fixme[6]: For 1st param expected `List[Parameter]` but got
-                #  `List[RangeParameter]`.
-                SumConstraint(parameters=parameters, is_upper_bound=True, bound=5)
-            ],
+            parameter_constraints=[ParameterConstraint(inequality="x + y <= 5")],
         )
         t = IntToFloat(search_space=constrained_int_search_space)
         self.assertEqual(t.rounding, "randomized")
@@ -288,11 +284,7 @@ def test_RoundingWithImpossiblyConstrainedIntRanges(self) -> None:
         ]
         constrained_int_search_space = SearchSpace(
             parameters=parameters,
-            parameter_constraints=[
-                # pyre-fixme[6]: For 1st param expected `List[Parameter]` but got
-                #  `List[RangeParameter]`.
-                SumConstraint(parameters=parameters, is_upper_bound=True, bound=3)
-            ],
+            parameter_constraints=[ParameterConstraint(inequality="x + y <= 3")],
         )
         t = IntToFloat(search_space=constrained_int_search_space)
         self.assertEqual(t.rounding, "randomized")

ax/core/__init__.py

@@ -31,7 +31,7 @@
     ParameterType,
     RangeParameter,
 )
-from ax.core.parameter_constraint import ParameterConstraint, SumConstraint
+from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.runner import Runner
 from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
@@ -61,6 +61,5 @@
     "Runner",
     "SearchSpace",
     "SimpleExperiment",
-    "SumConstraint",
     "Trial",
 ]

ax/core/parameter_constraint.py

@@ -11,7 +11,6 @@
 from typing import Sequence
 
 from ax.core.parameter import Parameter, RangeParameter
-from ax.core.types import ComparisonOp
 from ax.exceptions.core import UserInputError
 from ax.utils.common.base import SortableBase
 from ax.utils.common.string_utils import unsanitize_name
@@ -114,92 +113,6 @@ def _unique_id(self) -> str:
         return str(self)
 
 
-class SumConstraint(ParameterConstraint):
-    """Constraint on the sum of parameters being greater or less than a bound."""
-
-    def __init__(
-        self, parameters: list[Parameter], is_upper_bound: bool, bound: float
-    ) -> None:
-        """Initialize SumConstraint
-
-        Args:
-            parameters: List of parameters whose sum to constrain on.
-            is_upper_bound: Whether the bound is an upper or lower bound on the sum.
-            bound: The bound on the sum.
-        """
-        validate_constraint_parameters(parameters)
-
-        self._parameters = parameters
-        self._is_upper_bound: bool = is_upper_bound
-        self._parameter_names: list[str] = [parameter.name for parameter in parameters]
-        self._bound: float = self._inequality_weight * bound
-        self._constraint_dict: dict[str, float] = {
-            name: self._inequality_weight for name in self._parameter_names
-        }
-
-    @property
-    def parameters(self) -> list[Parameter]:
-        """Parameters."""
-        return self._parameters
-
-    @property
-    def constraint_dict(self) -> dict[str, float]:
-        """Weights on parameters for linear constraint representation."""
-        return self._constraint_dict
-
-    @property
-    def op(self) -> ComparisonOp:
-        """Whether the sum is constrained by a <= or >= inequality."""
-        return ComparisonOp.LEQ if self._is_upper_bound else ComparisonOp.GEQ
-
-    @property
-    def is_upper_bound(self) -> bool:
-        """Whether the bound is an upper or lower bound on the sum."""
-        return self._is_upper_bound
-
-    def clone(self) -> SumConstraint:
-        """Clone.
-
-        To use the same constraint, we need to reconstruct the original bound.
-        We do this by re-applying the original bound weighting.
-        """
-        return SumConstraint(
-            parameters=[p.clone() for p in self._parameters],
-            is_upper_bound=self._is_upper_bound,
-            bound=self._inequality_weight * self._bound,
-        )
-
-    def clone_with_transformed_parameters(
-        self, transformed_parameters: dict[str, Parameter]
-    ) -> SumConstraint:
-        """Clone, but replace parameters with transformed versions."""
-        return SumConstraint(
-            parameters=[transformed_parameters[p.name] for p in self._parameters],
-            is_upper_bound=self._is_upper_bound,
-            bound=self._inequality_weight * self._bound,
-        )
-
-    @property
-    def _inequality_weight(self) -> float:
-        """Multiplier of all terms in the inequality.
-
-        If the constraint is an upper bound, it is v1 + v2 ... v_n <= b
-        If the constraint is an lower bound, it is -v1 + -v2 ... -v_n <= -b
-        This property returns 1 or -1 depending on the scenario
-        """
-        return 1.0 if self._is_upper_bound else -1.0
-
-    def __repr__(self) -> str:
-        symbol = ">=" if self.op == ComparisonOp.GEQ else "<="
-        return (
-            "SumConstraint("
-            + " + ".join(self._parameter_names)
-            + " {} {})".format(
-                symbol, self._bound if self.op == ComparisonOp.LEQ else -self._bound
-            )
-        )
-
-
 def validate_constraint_parameters(parameters: Sequence[Parameter]) -> None:
     """Basic validation of parameters used in a constraint.
 

ax/core/search_space.py

@@ -28,7 +28,6 @@
 )
 from ax.core.parameter_constraint import (
     ParameterConstraint,
-    SumConstraint,
     validate_constraint_parameters,
 )
 from ax.core.types import TParameterization
@@ -190,22 +189,15 @@ def set_parameter_constraints(
         # Validate that all parameters in constraints are in search
         # space already.
         self._validate_parameter_constraints(parameter_constraints)
-        # Set the parameter on the constraint to be the parameter by
-        # 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, SumConstraint):
-                for idx, parameter in enumerate(constraint.parameters):
-                    constraint.parameters[idx] = self.parameters[parameter.name]
+        self._parameter_constraints: list[ParameterConstraint] = parameter_constraints
 
+        for constraint in self.parameter_constraints:
             validate_constraint_parameters(
                 parameters=[
                     self._parameters[name] for name in constraint.constraint_dict.keys()
                 ]
             )
 
-        self._parameter_constraints: list[ParameterConstraint] = parameter_constraints
-
     def add_parameters(
         self,
         parameters: Sequence[Parameter],
@@ -533,29 +525,17 @@ def _validate_parameter_constraints(
         self, parameter_constraints: list[ParameterConstraint]
     ) -> None:
         for constraint in parameter_constraints:
-            if isinstance(constraint, SumConstraint):
-                for parameter in constraint.parameters:
-                    if parameter.name not in self._parameters.keys():
-                        raise ValueError(
-                            f"`{parameter.name}` does not exist in search space."
-                        )
-                    if parameter != self._parameters[parameter.name]:
-                        raise ValueError(
-                            f"Parameter constraint's definition of '{parameter.name}' "
-                            "does not match the SearchSpace's definition"
-                        )
-            else:
-                for parameter_name in constraint.constraint_dict.keys():
-                    p = self._parameters.get(parameter_name)
-                    if p is None:
-                        raise ValueError(
-                            f"`{parameter_name}` does not exist in search space."
-                        )
-                    elif isinstance(p, DerivedParameter):
-                        raise ValueError(
-                            "Parameter constraints cannot be used with derived "
-                            "parameters."
-                        )
+            for parameter_name in constraint.constraint_dict.keys():
+                p = self._parameters.get(parameter_name)
+                if p is None:
+                    raise ValueError(
+                        f"`{parameter_name}` does not exist in search space."
+                    )
+                elif isinstance(p, DerivedParameter):
+                    raise ValueError(
+                        "Parameter constraints cannot be used with derived "
+                        "parameters."
+                    )
 
     def _validate_hierarchical_structure(self) -> None:
         """Validate the structure of this hierarchical search space, ensuring that all

ax/core/tests/test_parameter_constraint.py

@@ -6,12 +6,7 @@
 
 # pyre-strict
 
-from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
-from ax.core.parameter_constraint import (
-    ComparisonOp,
-    ParameterConstraint,
-    SumConstraint,
-)
+from ax.core.parameter_constraint import ParameterConstraint
 from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
 
@@ -127,68 +122,3 @@ def test_Sortable(self) -> None:
             inequality="2 * x - 3 * y <= 6.0",
         )
         self.assertTrue(constraint1 < constraint2)
-
-
-class SumConstraintTest(TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-        self.x = RangeParameter("x", ParameterType.INT, lower=-5, upper=5)
-        self.y = RangeParameter("y", ParameterType.INT, lower=-5, upper=5)
-        self.constraint1 = SumConstraint(
-            parameters=[self.x, self.y], is_upper_bound=True, bound=5
-        )
-        self.constraint2 = SumConstraint(
-            parameters=[self.x, self.y], is_upper_bound=False, bound=-5
-        )
-
-        self.constraint_repr1 = "SumConstraint(x + y <= 5.0)"
-        self.constraint_repr2 = "SumConstraint(x + y >= -5.0)"
-
-    def test_BadConstruct(self) -> None:
-        with self.assertRaises(ValueError):
-            SumConstraint(parameters=[self.x, self.x], is_upper_bound=False, bound=-5.0)
-        z = ChoiceParameter("z", ParameterType.STRING, ["a", "b", "c"])
-        with self.assertRaises(ValueError):
-            # pyre-fixme[16]: `SumConstraintTest` has no attribute `constraint`.
-            self.constraint = SumConstraint(
-                parameters=[self.x, z], is_upper_bound=False, bound=-5.0
-            )
-
-    def test_Properties(self) -> None:
-        self.assertEqual(self.constraint1.op, ComparisonOp.LEQ)
-        self.assertTrue(self.constraint1._is_upper_bound)
-
-        self.assertEqual(self.constraint2.op, ComparisonOp.GEQ)
-        self.assertFalse(self.constraint2._is_upper_bound)
-
-    def test_Repr(self) -> None:
-        self.assertEqual(str(self.constraint1), self.constraint_repr1)
-        self.assertEqual(str(self.constraint2), self.constraint_repr2)
-
-    def test_Validate(self) -> None:
-        self.assertTrue(self.constraint1.check({"x": 1, "y": 4}))
-        self.assertTrue(self.constraint1.check({"x": 4, "y": 1}))
-        self.assertFalse(self.constraint1.check({"x": 1, "y": 5}))
-
-        self.assertTrue(self.constraint2.check({"x": -4, "y": -1}))
-        self.assertTrue(self.constraint2.check({"x": -1, "y": -4}))
-        self.assertFalse(self.constraint2.check({"x": -5, "y": -1}))
-
-    def test_Clone(self) -> None:
-        constraint_clone = self.constraint1.clone()
-        self.assertEqual(self.constraint1.bound, constraint_clone.bound)
-
-        constraint_clone._bound = 7.0
-        self.assertNotEqual(self.constraint1.bound, constraint_clone.bound)
-
-        constraint_clone_2 = self.constraint2.clone()
-        self.assertEqual(self.constraint2.bound, constraint_clone_2.bound)
-
-    def test_CloneWithTransformedParameters(self) -> None:
-        constraint_clone = self.constraint1.clone_with_transformed_parameters(
-            transformed_parameters={p.name: p for p in self.constraint1.parameters}
-        )
-        self.assertEqual(self.constraint1.bound, constraint_clone.bound)
-
-        constraint_clone._bound = 7.0
-        self.assertNotEqual(self.constraint1.bound, constraint_clone.bound)

ax/core/tests/test_search_space.py

@@ -20,7 +20,7 @@
     ParameterType,
     RangeParameter,
 )
-from ax.core.parameter_constraint import ParameterConstraint, SumConstraint
+from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace, SearchSpaceDigest
 from ax.core.types import TParameterization
 from ax.exceptions.core import UserInputError
@@ -150,9 +150,7 @@ def test_Repr(self) -> None:
         self.assertEqual(str(self.ss1), self.ss1_repr)
 
     def test_Setter(self) -> None:
-        new_c = SumConstraint(
-            parameters=[self.a, self.b], is_upper_bound=True, bound=10
-        )
+        new_c = ParameterConstraint(inequality="a + b <= 10")
         self.ss2.add_parameter_constraints([new_c])
         self.assertEqual(len(self.ss2.parameter_constraints), 2)