Replace individual constraint methods by single get_constraints

miplearn/__init__.py

@@ -16,7 +16,6 @@ from .components.static_lazy import StaticLazyConstraintsComponent
 from .features import (
     Features,
     TrainingSample,
-    ConstraintFeatures,
     VariableFeatures,
     InstanceFeatures,
 )

miplearn/features.py

@@ -42,20 +42,20 @@ class VariableFeatures:
 
 
 @dataclass
-class ConstraintFeatures:
-    rhs: Optional[float] = None
-    lhs: Optional[Dict[str, float]] = None
-    sense: Optional[str] = None
-    category: Optional[Hashable] = None
+class Constraint:
+    rhs: float = 0.0
+    lhs: Dict[str, float] = lambda: {}  # type: ignore
+    sense: str = "<"
     user_features: Optional[List[float]] = None
     lazy: bool = False
+    category: Hashable = None
 
 
 @dataclass
 class Features:
     instance: Optional[InstanceFeatures] = None
     variables: Optional[Dict[str, VariableFeatures]] = None
-    constraints: Optional[Dict[str, ConstraintFeatures]] = None
+    constraints: Optional[Dict[str, Constraint]] = None
 
 
 class FeaturesExtractor:
@@ -106,10 +106,11 @@ class FeaturesExtractor:
     def _extract_constraints(
         self,
         instance: "Instance",
-    ) -> Dict[str, ConstraintFeatures]:
+    ) -> Dict[str, Constraint]:
         has_static_lazy = instance.has_static_lazy_constraints()
-        constraints: Dict[str, ConstraintFeatures] = {}
-        for cid in self.solver.get_constraint_ids():
+        constraints = self.solver.get_constraints()
+
+        for (cid, constr) in constraints.items():
             user_features = None
             category = instance.get_constraint_category(cid)
             if category is not None:
@@ -128,13 +129,8 @@ class FeaturesExtractor:
                     f"Constraint features must be a list of floats. "
                     f"Found {type(user_features[0]).__name__} instead for cid={cid}."
                 )
-            constraints[cid] = ConstraintFeatures(
-                rhs=self.solver.get_constraint_rhs(cid),
-                lhs=self.solver.get_constraint_lhs(cid),
-                sense=self.solver.get_constraint_sense(cid),
-                category=category,
-                user_features=user_features,
-            )
+            constraints[cid].category = category
+            constraints[cid].user_features = user_features
             if has_static_lazy:
                 constraints[cid].lazy = instance.is_constraint_lazy(cid)
         return constraints

miplearn/solvers/gurobi.py

@@ -10,6 +10,7 @@ from typing import List, Any, Dict, Optional, Hashable
 
 from overrides import overrides
 
+from miplearn.features import Constraint
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -25,7 +26,6 @@ from miplearn.types import (
     UserCutCallback,
     Solution,
     VariableName,
-    Constraint,
 )
 
 logger = logging.getLogger(__name__)
@@ -325,29 +325,7 @@ class GurobiSolver(InternalSolver):
             var.ub = value
 
     @overrides
-    def get_constraint_ids(self) -> List[str]:
-        assert self.model is not None
-        self._raise_if_callback()
-        self.model.update()
-        return [c.ConstrName for c in self.model.getConstrs()]
-
-    @overrides
-    def get_constraint_rhs(self, cid: str) -> float:
-        assert self.model is not None
-        return self.model.getConstrByName(cid).rhs
-
-    @overrides
-    def get_constraint_lhs(self, cid: str) -> Dict[str, float]:
-        assert self.model is not None
-        constr = self.model.getConstrByName(cid)
-        expr = self.model.getRow(constr)
-        lhs: Dict[str, float] = {}
-        for i in range(expr.size()):
-            lhs[expr.getVar(i).varName] = expr.getCoeff(i)
-        return lhs
-
-    @overrides
-    def extract_constraint(self, cid: str) -> Constraint:
+    def extract_constraint(self, cid: str) -> Any:
         self._raise_if_callback()
         assert self.model is not None
         constr = self.model.getConstrByName(cid)
@@ -358,7 +336,7 @@ class GurobiSolver(InternalSolver):
     @overrides
     def is_constraint_satisfied(
         self,
-        cobj: Constraint,
+        cobj: Any,
         tol: float = 1e-6,
     ) -> bool:
         lhs, sense, rhs, name = cobj
@@ -385,18 +363,6 @@ class GurobiSolver(InternalSolver):
         ineqs = [c for c in self.model.getConstrs() if c.sense != "="]
         return {c.ConstrName: c.Slack for c in ineqs}
 
-    @overrides
-    def set_constraint_sense(self, cid: str, sense: str) -> None:
-        assert self.model is not None
-        c = self.model.getConstrByName(cid)
-        c.Sense = sense
-
-    @overrides
-    def get_constraint_sense(self, cid: str) -> str:
-        assert self.model is not None
-        c = self.model.getConstrByName(cid)
-        return c.Sense
-
     @overrides
     def relax(self) -> None:
         assert self.model is not None
@@ -460,6 +426,26 @@ class GurobiSolver(InternalSolver):
             capacity=67.0,
         )
 
+    @overrides
+    def get_constraints(self) -> Dict[str, Constraint]:
+        assert self.model is not None
+        self._raise_if_callback()
+        self.model.update()
+        constraints: Dict[str, Constraint] = {}
+        for c in self.model.getConstrs():
+            expr = self.model.getRow(c)
+            lhs: Dict[str, float] = {}
+            for i in range(expr.size()):
+                lhs[expr.getVar(i).varName] = expr.getCoeff(i)
+            assert c.constrName not in constraints
+            constraints[c.constrName] = Constraint(
+                rhs=c.rhs,
+                lhs=lhs,
+                sense=c.sense,
+            )
+
+        return constraints
+
 
 class GurobiTestInstanceInfeasible(Instance):
     @overrides

miplearn/solvers/internal.py

@@ -6,6 +6,9 @@ import logging
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
 
+from overrides import EnforceOverrides
+
+from miplearn.features import Constraint
 from miplearn.instance.base import Instance
 from miplearn.types import (
     LPSolveStats,
@@ -13,7 +16,6 @@ from miplearn.types import (
     LazyCallback,
     MIPSolveStats,
     BranchPriorities,
-    Constraint,
     UserCutCallback,
     Solution,
     VariableName,
@@ -151,32 +153,11 @@ class InternalSolver(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def get_constraint_ids(self) -> List[str]:
-        """
-        Returns a list of ids which uniquely identify each constraint in the model.
-        """
-        pass
-
-    @abstractmethod
-    def get_constraint_rhs(self, cid: str) -> float:
-        """
-        Returns the right-hand side of a given constraint.
-        """
+    def get_constraints(self) -> Dict[str, Constraint]:
         pass
 
     @abstractmethod
-    def get_constraint_lhs(self, cid: str) -> Dict[str, float]:
-        """
-        Returns a list of tuples encoding the left-hand side of the constraint.
-
-        The first element of the tuple is the name of the variable and the second
-        element is the coefficient. For example, the left-hand side of "2 x1 + x2 <= 3"
-        is encoded as [{"x1": 2, "x2": 1}].
-        """
-        pass
-
-    @abstractmethod
-    def add_constraint(self, cobj: Constraint, name: str = "") -> None:
+    def add_constraint(self, cobj: Any, name: str = "") -> None:
         """
         Adds a single constraint to the model.
         """
@@ -190,7 +171,7 @@ class InternalSolver(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def extract_constraint(self, cid: str) -> Constraint:
+    def extract_constraint(self, cid: str) -> Any:
         """
         Removes a given constraint from the model and returns an object `cobj` which
         can be used to verify if the removed constraint is still satisfied by
@@ -200,38 +181,12 @@ class InternalSolver(ABC):
         pass
 
     @abstractmethod
-    def is_constraint_satisfied(self, cobj: Constraint, tol: float = 1e-6) -> bool:
+    def is_constraint_satisfied(self, cobj: Any, tol: float = 1e-6) -> bool:
         """
         Returns True if the current solution satisfies the given constraint.
         """
         pass
 
-    @abstractmethod
-    def set_constraint_sense(self, cid: str, sense: str) -> None:
-        """
-        Modifies the sense of a given constraint.
-
-        Parameters
-        ----------
-        cid: str
-            The name of the constraint.
-        sense: str
-            The new sense (either "<", ">" or "=").
-        """
-        pass
-
-    @abstractmethod
-    def get_constraint_sense(self, cid: str) -> str:
-        """
-        Returns the sense of a given constraint (either "<", ">" or "=").
-
-        Parameters
-        ----------
-        cid: str
-            The name of the constraint.
-        """
-        pass
-
     @abstractmethod
     def relax(self) -> None:
         """

miplearn/solvers/pyomo/base.py

@@ -11,7 +11,10 @@ from typing import Any, List, Dict, Optional, Hashable
 import pyomo
 from overrides import overrides
 from pyomo import environ as pe
-from pyomo.core import Var, Constraint
+from pyomo.core import Var
+from pyomo.core.base import _GeneralVarData
+from pyomo.core.base.constraint import SimpleConstraint
+from pyomo.core.expr.numeric_expr import SumExpression, MonomialTermExpression
 from pyomo.opt import TerminationCondition
 from pyomo.opt.base.solvers import SolverFactory
 
@@ -23,6 +26,7 @@ from miplearn.solvers.internal import (
     IterationCallback,
     LazyCallback,
     MIPSolveStats,
+    Constraint,
 )
 from miplearn.types import (
     SolverParams,
@@ -213,7 +217,7 @@ class BasePyomoSolver(InternalSolver):
     def _update_constrs(self) -> None:
         assert self.model is not None
         self._cname_to_constr = {}
-        for constr in self.model.component_objects(Constraint):
+        for constr in self.model.component_objects(pyomo.core.Constraint):
             if isinstance(constr, pe.ConstraintList):
                 for idx in constr:
                     self._cname_to_constr[f"{constr.name}[{idx}]"] = constr[idx]
@@ -262,10 +266,6 @@ class BasePyomoSolver(InternalSolver):
             return None
         return int(value)
 
-    @overrides
-    def get_constraint_ids(self) -> List[str]:
-        return list(self._cname_to_constr.keys())
-
     def _get_warm_start_regexp(self) -> Optional[str]:
         return None
 
@@ -290,37 +290,6 @@ class BasePyomoSolver(InternalSolver):
             result[cname] = cobj.slack()
         return result
 
-    @overrides
-    def get_constraint_sense(self, cid: str) -> str:
-        cobj = self._cname_to_constr[cid]
-        has_ub = cobj.has_ub()
-        has_lb = cobj.has_lb()
-        assert (
-            (not has_lb) or (not has_ub) or cobj.upper() == cobj.lower()
-        ), "range constraints not supported"
-        if has_lb:
-            return ">"
-        elif has_ub:
-            return "<"
-        else:
-            return "="
-
-    @overrides
-    def get_constraint_rhs(self, cid: str) -> float:
-        cobj = self._cname_to_constr[cid]
-        if cobj.has_ub:
-            return cobj.upper()
-        else:
-            return cobj.lower()
-
-    @overrides
-    def get_constraint_lhs(self, cid: str) -> Dict[str, float]:
-        return {}
-
-    @overrides
-    def set_constraint_sense(self, cid: str, sense: str) -> None:
-        raise NotImplementedError()
-
     @overrides
     def extract_constraint(self, cid: str) -> Constraint:
         raise NotImplementedError()
@@ -357,6 +326,63 @@ class BasePyomoSolver(InternalSolver):
             capacity=67.0,
         )
 
+    @overrides
+    def get_constraints(self) -> Dict[str, Constraint]:
+        assert self.model is not None
+
+        def _get(c: pyomo.core.Constraint, name: str) -> Constraint:
+            # Extract RHS and sense
+            has_ub = c.has_ub()
+            has_lb = c.has_lb()
+            assert (
+                (not has_lb) or (not has_ub) or c.upper() == c.lower()
+            ), "range constraints not supported"
+            if has_lb:
+                sense = ">"
+                rhs = c.lower()
+            elif has_ub:
+                sense = "<"
+                rhs = c.upper()
+            else:
+                sense = "="
+                rhs = c.upper()
+
+            # Extract LHS
+            lhs = {}
+            if isinstance(c.body, SumExpression):
+                for term in c.body._args_:
+                    if isinstance(term, MonomialTermExpression):
+                        lhs[term._args_[1].name] = term._args_[0]
+                    elif isinstance(term, _GeneralVarData):
+                        lhs[term.name] = 1.0
+                    else:
+                        raise Exception(f"Unknown term type: {term.__class__.__name__}")
+            elif isinstance(c.body, _GeneralVarData):
+                lhs[c.body.name] = 1.0
+            else:
+                raise Exception(f"Unknown expression type: {c.body.__class__.__name__}")
+
+            # Build constraint
+            return Constraint(
+                lhs=lhs,
+                rhs=rhs,
+                sense=sense,
+            )
+
+        constraints = {}
+        for constr in self.model.component_objects(pyomo.core.Constraint):
+            if isinstance(constr, pe.ConstraintList):
+                for idx in constr:
+                    name = f"{constr.name}[{idx}]"
+                    assert name not in constraints
+                    constraints[name] = _get(constr[idx], name=name)
+            else:
+                name = constr.name
+                assert name not in constraints
+                constraints[name] = _get(constr, name=name)
+
+        return constraints
+
 
 class PyomoTestInstanceInfeasible(Instance):
     @overrides

miplearn/solvers/tests/__init__.py

@@ -4,8 +4,10 @@
 
 from typing import Any
 
+from miplearn.features import Constraint
 from miplearn.solvers.internal import InternalSolver
 
+
 # NOTE:
 # This file is in the main source folder, so that it can be called from Julia.
 
@@ -66,6 +68,22 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
     assert_equals(solution["x[2]"], 1.0)
     assert_equals(solution["x[3]"], 1.0)
 
+    assert_equals(
+        solver.get_constraints(),
+        {
+            "eq_capacity": Constraint(
+                lhs={
+                    "x[0]": 23.0,
+                    "x[1]": 26.0,
+                    "x[2]": 20.0,
+                    "x[3]": 18.0,
+                },
+                rhs=67.0,
+                sense="<",
+            ),
+        },
+    )
+
     # assert_equals(solver.get_constraint_ids(), ["eq_capacity"])
     # assert_equals(
     #     solver.get_constraint_rhs("eq_capacity"),
@@ -96,16 +114,34 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
     assert cut is not None
     solver.add_constraint(cut, name="cut")
 
-    # New constraint should affect solution and should be listed in
-    # constraint ids
-    assert solver.get_constraint_ids() == ["eq_capacity", "cut"]
+    # New constraint should be listed
+    assert_equals(
+        solver.get_constraints(),
+        {
+            "eq_capacity": Constraint(
+                lhs={
+                    "x[0]": 23.0,
+                    "x[1]": 26.0,
+                    "x[2]": 20.0,
+                    "x[3]": 18.0,
+                },
+                rhs=67.0,
+                sense="<",
+            ),
+            "cut": Constraint(
+                lhs={
+                    "x[0]": 1.0,
+                },
+                rhs=0.0,
+                sense="<",
+            ),
+        },
+    )
+
+    # New constraint should affect the solution
     stats = solver.solve()
     assert stats["Lower bound"] == 1030.0
 
-    assert solver.get_sense() == "max"
-    assert solver.get_constraint_sense("cut") == "<"
-    assert solver.get_constraint_sense("eq_capacity") == "<"
-
     # Verify slacks
     assert solver.get_inequality_slacks() == {
         "cut": 0.0,

miplearn/types.py

@@ -12,7 +12,6 @@ if TYPE_CHECKING:
 
 BranchPriorities = Dict[str, Optional[float]]
 Category = Hashable
-Constraint = Any
 IterationCallback = Callable[[], bool]
 LazyCallback = Callable[[Any, Any], None]
 SolverParams = Dict[str, Any]

tests/components/test_static_lazy.py

@@ -14,8 +14,8 @@ from miplearn.components.static_lazy import StaticLazyConstraintsComponent
 from miplearn.features import (
     TrainingSample,
     InstanceFeatures,
-    ConstraintFeatures,
     Features,
+    Constraint,
 )
 from miplearn.instance.base import Instance
 from miplearn.solvers.internal import InternalSolver
@@ -48,27 +48,27 @@ def features() -> Features:
             lazy_constraint_count=4,
         ),
         constraints={
-            "c1": ConstraintFeatures(
+            "c1": Constraint(
                 category="type-a",
                 user_features=[1.0, 1.0],
                 lazy=True,
             ),
-            "c2": ConstraintFeatures(
+            "c2": Constraint(
                 category="type-a",
                 user_features=[1.0, 2.0],
                 lazy=True,
             ),
-            "c3": ConstraintFeatures(
+            "c3": Constraint(
                 category="type-a",
                 user_features=[1.0, 3.0],
                 lazy=True,
             ),
-            "c4": ConstraintFeatures(
+            "c4": Constraint(
                 category="type-b",
                 user_features=[1.0, 4.0, 0.0],
                 lazy=True,
             ),
-            "c5": ConstraintFeatures(
+            "c5": Constraint(
                 category="type-b",
                 user_features=[1.0, 5.0, 0.0],
                 lazy=False,

tests/test_features.py

@@ -6,7 +6,7 @@ from miplearn.features import (
     FeaturesExtractor,
     InstanceFeatures,
     VariableFeatures,
-    ConstraintFeatures,
+    Constraint,
 )
 from miplearn.solvers.gurobi import GurobiSolver
 
@@ -37,7 +37,7 @@ def test_knapsack() -> None:
             ),
         }
         assert instance.features.constraints == {
-            "eq_capacity": ConstraintFeatures(
+            "eq_capacity": Constraint(
                 lhs={
                     "x[0]": 23.0,
                     "x[1]": 26.0,