Extract more features to ConstraintFeatures

miplearn/features.py

@@ -148,6 +148,7 @@ class Constraint:
 class Features:
     instance: Optional[InstanceFeatures] = None
     variables: Optional[VariableFeatures] = None
+    constraints: Optional[ConstraintFeatures] = None
     constraints_old: Optional[Dict[str, Constraint]] = None
     lp_solve: Optional["LPSolveStats"] = None
     mip_solve: Optional["MIPSolveStats"] = None
@@ -179,6 +180,10 @@ class FeaturesExtractor:
             with_static=with_static,
             with_sa=self.with_sa,
         )
+        features.constraints = solver.get_constraints(
+            with_static=with_static,
+            with_sa=self.with_sa,
+        )
         features.constraints_old = solver.get_constraints_old(
             with_static=with_static,
         )

miplearn/solvers/gurobi.py

@@ -149,29 +149,40 @@ class GurobiSolver(InternalSolver):
     def get_constraint_attrs(self) -> List[str]:
         return [
             "basis_status",
-            "category",
-            "dual_value",
+            "categories",
+            "dual_values",
             "lazy",
             "lhs",
+            "names",
             "rhs",
             "sa_rhs_down",
             "sa_rhs_up",
-            "sense",
-            "slack",
+            "senses",
+            "slacks",
             "user_features",
         ]
 
     @overrides
-    def get_constraints(self, with_static: bool = True) -> ConstraintFeatures:
+    def get_constraints(
+        self,
+        with_static: bool = True,
+        with_sa: bool = True,
+    ) -> ConstraintFeatures:
         model = self.model
         assert model is not None
         assert model.numVars == len(self._gp_vars)
 
+        def _parse_gurobi_cbasis(v: int) -> str:
+            if v == 0:
+                return "B"
+            if v == -1:
+                return "N"
+            raise Exception(f"unknown cbasis: {v}")
+
         gp_constrs = model.getConstrs()
         constr_names = tuple(model.getAttr("constrName", gp_constrs))
-        rhs = None
-        senses = None
-        lhs = None
+        rhs, lhs, senses, slacks, basis_status = None, None, None, None, None
+        dual_value, basis_status, sa_rhs_up, sa_rhs_down = None, None, None, None
 
         if with_static:
             rhs = tuple(model.getAttr("rhs", gp_constrs))
@@ -185,11 +196,31 @@ class GurobiSolver(InternalSolver):
                 )
             lhs = tuple(lhs_l)
 
+        if self._has_lp_solution:
+            dual_value = tuple(model.getAttr("pi", gp_constrs))
+            basis_status = tuple(
+                map(
+                    _parse_gurobi_cbasis,
+                    model.getAttr("cbasis", gp_constrs),
+                )
+            )
+            if with_sa:
+                sa_rhs_up = tuple(model.getAttr("saRhsUp", gp_constrs))
+                sa_rhs_down = tuple(model.getAttr("saRhsLow", gp_constrs))
+
+        if self._has_lp_solution or self._has_mip_solution:
+            slacks = tuple(model.getAttr("slack", gp_constrs))
+
         return ConstraintFeatures(
+            basis_status=basis_status,
+            dual_values=dual_value,
+            lhs=lhs,
             names=constr_names,
             rhs=rhs,
+            sa_rhs_down=sa_rhs_down,
+            sa_rhs_up=sa_rhs_up,
             senses=senses,
-            lhs=lhs,
+            slacks=slacks,
         )
 
     @overrides

miplearn/solvers/internal.py

@@ -169,7 +169,11 @@ class InternalSolver(ABC, EnforceOverrides):
         raise NotImplementedError()
 
     @abstractmethod
-    def get_constraints(self, with_static: bool = True) -> ConstraintFeatures:
+    def get_constraints(
+        self,
+        with_static: bool = True,
+        with_sa: bool = True,
+    ) -> ConstraintFeatures:
         pass
 
     @abstractmethod

miplearn/solvers/pyomo/base.py

@@ -128,15 +128,23 @@ class BasePyomoSolver(InternalSolver):
             self._pyomo_solver.update_var(var)
 
     @overrides
-    def get_constraints(self, with_static: bool = True) -> ConstraintFeatures:
-        assert self.model is not None
+    def get_constraints(
+        self,
+        with_static: bool = True,
+        with_sa: bool = True,
+    ) -> ConstraintFeatures:
+        model = self.model
+        assert model is not None
 
         names: List[str] = []
         rhs: List[float] = []
         lhs: List[Tuple[Tuple[str, float], ...]] = []
         senses: List[str] = []
+        dual_values: List[float] = []
+        slacks: List[float] = []
 
         def _parse_constraint(c: pe.Constraint) -> None:
+            assert model is not None
             if with_static:
                 # Extract RHS and sense
                 has_ub = c.has_ub()
@@ -175,7 +183,15 @@ class BasePyomoSolver(InternalSolver):
                     )
                 lhs.append(tuple(lhsc))
 
-        for constr in self.model.component_objects(pyomo.core.Constraint):
+            # Extract dual values
+            if self._has_lp_solution:
+                dual_values.append(model.dual[c])
+
+            # Extract slacks
+            if self._has_mip_solution or self._has_lp_solution:
+                slacks.append(model.slack[c])
+
+        for constr in model.component_objects(pyomo.core.Constraint):
             if isinstance(constr, pe.ConstraintList):
                 for idx in constr:
                     names.append(f"{constr.name}[{idx}]")
@@ -185,16 +201,23 @@ class BasePyomoSolver(InternalSolver):
                 _parse_constraint(constr)
 
         rhs_t, lhs_t, senses_t = None, None, None
+        slacks_t, dual_values_t = None, None
         if with_static:
             rhs_t = tuple(rhs)
             lhs_t = tuple(lhs)
             senses_t = tuple(senses)
+        if self._has_lp_solution:
+            dual_values_t = tuple(dual_values)
+        if self._has_lp_solution or self._has_mip_solution:
+            slacks_t = tuple(slacks)
 
         return ConstraintFeatures(
             names=tuple(names),
             rhs=rhs_t,
             senses=senses_t,
             lhs=lhs_t,
+            slacks=slacks_t,
+            dual_values=dual_values_t,
         )
 
     @overrides
@@ -224,12 +247,12 @@ class BasePyomoSolver(InternalSolver):
     @overrides
     def get_constraint_attrs(self) -> List[str]:
         return [
-            "dual_value",
-            "lazy",
+            "dual_values",
             "lhs",
+            "names",
             "rhs",
-            "sense",
-            "slack",
+            "senses",
+            "slacks",
         ]
 
     @overrides

miplearn/solvers/tests/__init__.py

@@ -30,16 +30,12 @@ def _round(obj: Any) -> Any:
         return tuple([_round(v) for v in obj])
     if isinstance(obj, list):
         return [_round(v) for v in obj]
+    if isinstance(obj, dict):
+        return {key: _round(value) for (key, value) in obj.items()}
     if isinstance(obj, VariableFeatures):
-        obj.reduced_costs = _round(obj.reduced_costs)
-        obj.sa_obj_up = _round(obj.sa_obj_up)
-        obj.sa_obj_down = _round(obj.sa_obj_down)
-        obj.sa_lb_up = _round(obj.sa_lb_up)
-        obj.sa_lb_down = _round(obj.sa_lb_down)
-        obj.sa_ub_up = _round(obj.sa_ub_up)
-        obj.sa_ub_down = _round(obj.sa_ub_down)
-        obj.values = _round(obj.values)
-        obj.alvarez_2017 = _round(obj.alvarez_2017)
+        obj.__dict__ = _round(obj.__dict__)
+    if isinstance(obj, ConstraintFeatures):
+        obj.__dict__ = _round(obj.__dict__)
     return obj
 
 
@@ -110,18 +106,6 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
         ),
     )
 
-    assert_equals(
-        _round_constraints(solver.get_constraints_old()),
-        {
-            "eq_capacity": Constraint(
-                lazy=False,
-                lhs={"x[0]": 23.0, "x[1]": 26.0, "x[2]": 20.0, "x[3]": 18.0, "z": -1.0},
-                rhs=0.0,
-                sense="=",
-            )
-        },
-    )
-
     # Solve linear programming relaxation
     lp_stats = solver.solve_lp()
     assert not solver.is_infeasible()
@@ -154,28 +138,17 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # Fetch constraints (after-lp)
     assert_equals(
-        _round_constraints(solver.get_constraints_old()),
-        _remove_unsupported_constr_attrs(
-            solver,
-            {
-                "eq_capacity": Constraint(
-                    basis_status="N",
-                    dual_value=13.538462,
-                    lazy=False,
-                    lhs={
-                        "x[0]": 23.0,
-                        "x[1]": 26.0,
-                        "x[2]": 20.0,
-                        "x[3]": 18.0,
-                        "z": -1.0,
-                    },
-                    rhs=0.0,
-                    sa_rhs_down=-24.0,
-                    sa_rhs_up=1.9999999999999987,
-                    sense="=",
-                    slack=0.0,
-                )
-            },
+        _round(solver.get_constraints(with_static=False)),
+        _filter_attrs(
+            solver.get_constraint_attrs(),
+            ConstraintFeatures(
+                basis_status=("N",),
+                dual_values=(13.538462,),
+                names=("eq_capacity",),
+                sa_rhs_down=(-24.0,),
+                sa_rhs_up=(2.0,),
+                slacks=(0.0,),
+            ),
         ),
     )