Create ConstraintFeatures

miplearn/features.py

@@ -76,6 +76,38 @@ class VariableFeatures:
         return features
 
 
+@dataclass
+class ConstraintFeatures:
+    basis_status: Optional[Tuple[str, ...]] = None
+    categories: Optional[Tuple[Optional[Hashable], ...]] = None
+    dual_values: Optional[Tuple[float, ...]] = None
+    names: Optional[Tuple[str, ...]] = None
+    lazy: Optional[Tuple[bool, ...]] = None
+    lhs: Optional[Tuple[Tuple[Tuple[str, float], ...], ...]] = None
+    rhs: Optional[Tuple[float, ...]] = None
+    sa_rhs_down: Optional[Tuple[float, ...]] = None
+    sa_rhs_up: Optional[Tuple[float, ...]] = None
+    senses: Optional[Tuple[str, ...]] = None
+    slacks: Optional[Tuple[float, ...]] = None
+    user_features: Optional[Tuple[Optional[Tuple[float, ...]], ...]] = None
+
+    def to_list(self, index: int) -> List[float]:
+        features: List[float] = []
+        for attr in [
+            "dual_values",
+            "rhs",
+            "slacks",
+        ]:
+            if getattr(self, attr) is not None:
+                features.append(getattr(self, attr)[index])
+        for attr in ["user_features"]:
+            if getattr(self, attr) is not None:
+                if getattr(self, attr)[index] is not None:
+                    features.extend(getattr(self, attr)[index])
+        _clip(features)
+        return features
+
+
 @dataclass
 class Constraint:
     basis_status: Optional[str] = None
@@ -147,7 +179,7 @@ class FeaturesExtractor:
             with_static=with_static,
             with_sa=self.with_sa,
         )
-        features.constraints_old = solver.get_constraints(
+        features.constraints_old = solver.get_constraints_old(
             with_static=with_static,
         )
         if with_static:

miplearn/solvers/gurobi.py

@@ -10,7 +10,7 @@ from typing import List, Any, Dict, Optional, Hashable, Tuple, TYPE_CHECKING
 
 from overrides import overrides
 
-from miplearn.features import Constraint, VariableFeatures
+from miplearn.features import Constraint, VariableFeatures, ConstraintFeatures
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -162,7 +162,38 @@ class GurobiSolver(InternalSolver):
         ]
 
     @overrides
-    def get_constraints(self, with_static: bool = True) -> Dict[str, Constraint]:
+    def get_constraints(self, with_static: bool = True) -> ConstraintFeatures:
+        model = self.model
+        assert model is not None
+        assert model.numVars == len(self._gp_vars)
+
+        gp_constrs = model.getConstrs()
+        constr_names = tuple(model.getAttr("constrName", gp_constrs))
+        rhs = None
+        senses = None
+        lhs = None
+
+        if with_static:
+            rhs = tuple(model.getAttr("rhs", gp_constrs))
+            senses = tuple(model.getAttr("sense", gp_constrs))
+            lhs_l: List = [None for _ in gp_constrs]
+            for (i, gp_constr) in enumerate(gp_constrs):
+                expr = model.getRow(gp_constr)
+                lhs_l[i] = tuple(
+                    (self._var_names[expr.getVar(j).index], expr.getCoeff(j))
+                    for j in range(expr.size())
+                )
+            lhs = tuple(lhs_l)
+
+        return ConstraintFeatures(
+            names=constr_names,
+            rhs=rhs,
+            senses=senses,
+            lhs=lhs,
+        )
+
+    @overrides
+    def get_constraints_old(self, with_static: bool = True) -> Dict[str, Constraint]:
         model = self.model
         assert model is not None
         self._raise_if_callback()
@@ -255,22 +286,6 @@ class GurobiSolver(InternalSolver):
     @overrides
     def get_variable_attrs(self) -> List[str]:
         return [
-            "basis_status",
-            "category",
-            "lower_bound",
-            "obj_coeff",
-            "reduced_cost",
-            "sa_lb_down",
-            "sa_lb_up",
-            "sa_obj_down",
-            "sa_obj_up",
-            "sa_ub_down",
-            "sa_ub_up",
-            "type",
-            "upper_bound",
-            "user_features",
-            "value",
-            # new attributes
             "names",
             "basis_status",
             "categories",

miplearn/solvers/internal.py

@@ -9,7 +9,7 @@ from typing import Any, Dict, List, Optional
 
 from overrides import EnforceOverrides
 
-from miplearn.features import Constraint, VariableFeatures
+from miplearn.features import Constraint, VariableFeatures, ConstraintFeatures
 from miplearn.instance.base import Instance
 from miplearn.types import (
     IterationCallback,
@@ -169,7 +169,11 @@ class InternalSolver(ABC, EnforceOverrides):
         raise NotImplementedError()
 
     @abstractmethod
-    def get_constraints(self, with_static: bool = True) -> Dict[str, Constraint]:
+    def get_constraints(self, with_static: bool = True) -> ConstraintFeatures:
+        pass
+
+    @abstractmethod
+    def get_constraints_old(self, with_static: bool = True) -> Dict[str, Constraint]:
         pass
 
     @abstractmethod

miplearn/solvers/pyomo/base.py

@@ -19,7 +19,7 @@ from pyomo.core.expr.numeric_expr import SumExpression, MonomialTermExpression
 from pyomo.opt import TerminationCondition
 from pyomo.opt.base.solvers import SolverFactory
 
-from miplearn.features import VariableFeatures
+from miplearn.features import VariableFeatures, ConstraintFeatures
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -128,7 +128,77 @@ class BasePyomoSolver(InternalSolver):
             self._pyomo_solver.update_var(var)
 
     @overrides
-    def get_constraints(self, with_static: bool = True) -> Dict[str, Constraint]:
+    def get_constraints(self, with_static: bool = True) -> ConstraintFeatures:
+        assert self.model is not None
+
+        names: List[str] = []
+        rhs: List[float] = []
+        lhs: List[Tuple[Tuple[str, float], ...]] = []
+        senses: List[str] = []
+
+        def _parse_constraint(c: pe.Constraint) -> None:
+            if with_static:
+                # 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 not has_ub:
+                    senses.append(">")
+                    rhs.append(float(c.lower()))
+                elif not has_lb:
+                    senses.append("<")
+                    rhs.append(float(c.upper()))
+                else:
+                    senses.append("=")
+                    rhs.append(float(c.upper()))
+
+                # Extract LHS
+                lhsc = []
+                expr = c.body
+                if isinstance(expr, SumExpression):
+                    for term in expr._args_:
+                        if isinstance(term, MonomialTermExpression):
+                            lhsc.append((term._args_[1].name, float(term._args_[0])))
+                        elif isinstance(term, _GeneralVarData):
+                            lhsc.append((term.name, 1.0))
+                        else:
+                            raise Exception(
+                                f"Unknown term type: {term.__class__.__name__}"
+                            )
+                elif isinstance(expr, _GeneralVarData):
+                    lhsc.append((expr.name, 1.0))
+                else:
+                    raise Exception(
+                        f"Unknown expression type: {expr.__class__.__name__}"
+                    )
+                lhs.append(tuple(lhsc))
+
+        for constr in self.model.component_objects(pyomo.core.Constraint):
+            if isinstance(constr, pe.ConstraintList):
+                for idx in constr:
+                    names.append(f"{constr.name}[{idx}]")
+                    _parse_constraint(constr[idx])
+            else:
+                names.append(constr.name)
+                _parse_constraint(constr)
+
+        rhs_t, lhs_t, senses_t = None, None, None
+        if with_static:
+            rhs_t = tuple(rhs)
+            lhs_t = tuple(lhs)
+            senses_t = tuple(senses)
+
+        return ConstraintFeatures(
+            names=tuple(names),
+            rhs=rhs_t,
+            senses=senses_t,
+            lhs=lhs_t,
+        )
+
+    @overrides
+    def get_constraints_old(self, with_static: bool = True) -> Dict[str, Constraint]:
         assert self.model is not None
 
         constraints = {}
@@ -267,20 +337,6 @@ class BasePyomoSolver(InternalSolver):
     @overrides
     def get_variable_attrs(self) -> List[str]:
         return [
-            # "basis_status",
-            "lower_bound",
-            "obj_coeff",
-            "reduced_cost",
-            # "sa_lb_down",
-            # "sa_lb_up",
-            # "sa_obj_down",
-            # "sa_obj_up",
-            # "sa_ub_down",
-            # "sa_ub_up",
-            "type",
-            "upper_bound",
-            "value",
-            # new attributes
             "names",
             # "basis_status",
             "categories",

miplearn/solvers/tests/__init__.py

@@ -4,7 +4,7 @@
 
 from typing import Any, Dict, List
 
-from miplearn.features import Constraint, VariableFeatures
+from miplearn.features import Constraint, VariableFeatures, ConstraintFeatures
 from miplearn.solvers.internal import InternalSolver
 
 inf = float("inf")
@@ -93,7 +93,25 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # Fetch constraints (after-load)
     assert_equals(
-        _round_constraints(solver.get_constraints()),
+        solver.get_constraints(),
+        ConstraintFeatures(
+            names=("eq_capacity",),
+            rhs=(0.0,),
+            lhs=(
+                (
+                    ("x[0]", 23.0),
+                    ("x[1]", 26.0),
+                    ("x[2]", 20.0),
+                    ("x[3]", 18.0),
+                    ("z", -1.0),
+                ),
+            ),
+            senses=("=",),
+        ),
+    )
+
+    assert_equals(
+        _round_constraints(solver.get_constraints_old()),
         {
             "eq_capacity": Constraint(
                 lazy=False,
@@ -136,7 +154,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # Fetch constraints (after-lp)
     assert_equals(
-        _round_constraints(solver.get_constraints()),
+        _round_constraints(solver.get_constraints_old()),
         _remove_unsupported_constr_attrs(
             solver,
             {
@@ -192,7 +210,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # Fetch constraints (after-mip)
     assert_equals(
-        _round_constraints(solver.get_constraints(with_static=False)),
+        _round_constraints(solver.get_constraints_old(with_static=False)),
         {"eq_capacity": Constraint(slack=0.0)},
     )
 
@@ -204,7 +222,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
     # also clear the current solution.
     solver.add_constraint(cut, "cut")
     assert_equals(
-        _round_constraints(solver.get_constraints()),
+        _round_constraints(solver.get_constraints_old()),
         {
             "eq_capacity": Constraint(
                 lazy=False,