PyomoSolver: Implement missing constraint methods

miplearn/solvers/internal.py

@@ -6,8 +6,6 @@ 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 (

miplearn/solvers/pyomo/base.py

@@ -13,7 +13,7 @@ from overrides import overrides
 from pyomo import environ as pe
 from pyomo.core import Var
 from pyomo.core.base import _GeneralVarData
-from pyomo.core.base.constraint import SimpleConstraint
+from pyomo.core.base.constraint import SimpleConstraint, ConstraintList
 from pyomo.core.expr.numeric_expr import SumExpression, MonomialTermExpression
 from pyomo.opt import TerminationCondition
 from pyomo.opt.base.solvers import SolverFactory
@@ -186,6 +186,7 @@ class BasePyomoSolver(InternalSolver):
         assert isinstance(model, pe.ConcreteModel)
         self.instance = instance
         self.model = model
+        self.model.extra_constraints = ConstraintList()
         self._pyomo_solver.set_instance(model)
         self._update_obj()
         self._update_vars()
@@ -234,8 +235,23 @@ class BasePyomoSolver(InternalSolver):
             self._pyomo_solver.update_var(var)
 
     @overrides
-    def add_constraint(self, constraint: Any, name: str = "") -> Any:
+    def add_constraint(self, cobj: Any, name: str = "") -> Any:
-        self._pyomo_solver.add_constraint(constraint)
+        assert self.model is not None
+        if isinstance(cobj, Constraint):
+            lhs = 0.0
+            for (varname, coeff) in cobj.lhs.items():
+                var = self._varname_to_var[varname]
+                lhs += var * coeff
+            if cobj.sense == "=":
+                expr = lhs == cobj.rhs
+            elif cobj.sense == "<":
+                expr = lhs <= cobj.rhs
+            else:
+                expr = lhs >= cobj.rhs
+            cl = self.model.extra_constraints
+            self._pyomo_solver.add_constraint(cl.add(expr))
+        else:
+            self._pyomo_solver.add_constraint(cobj)
         self._update_constrs()
 
     @staticmethod
@@ -291,12 +307,25 @@ class BasePyomoSolver(InternalSolver):
         return result
 
     @overrides
-    def extract_constraint(self, cid: str) -> Constraint:
+    def extract_constraint(self, cid: str) -> Any:
-        raise NotImplementedError()
+        cobj = self._cname_to_constr[cid]
+        constr = self._parse_pyomo_constraint(cobj)
+        self._pyomo_solver.remove_constraint(cobj)
+        return constr
 
     @overrides
-    def is_constraint_satisfied(self, cobj: Constraint, tol: float = 1e-6) -> bool:
+    def is_constraint_satisfied(self, cobj: Any, tol: float = 1e-6) -> bool:
-        raise NotImplementedError()
+        assert isinstance(cobj, Constraint)
+        lhs_value = 0.0
+        for (varname, coeff) in cobj.lhs.items():
+            var = self._varname_to_var[varname]
+            lhs_value += var.value * coeff
+        if cobj.sense == "=":
+            return (lhs_value <= cobj.rhs + tol) and (lhs_value >= cobj.rhs - tol)
+        elif cobj.sense == "<":
+            return lhs_value <= cobj.rhs + tol
+        else:
+            return lhs_value >= cobj.rhs - tol
 
     @overrides
     def is_infeasible(self) -> bool:
@@ -330,59 +359,60 @@ class BasePyomoSolver(InternalSolver):
     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)
+                    constraints[name] = self._parse_pyomo_constraint(constr[idx])
             else:
                 name = constr.name
                 assert name not in constraints
-                constraints[name] = _get(constr, name=name)
+                constraints[name] = self._parse_pyomo_constraint(constr)
 
         return constraints
 
+    @staticmethod
+    def _parse_pyomo_constraint(c: pyomo.core.Constraint) -> 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,
+        )
+
 
 class PyomoTestInstanceInfeasible(Instance):
     @overrides

miplearn/solvers/tests/__init__.py

@@ -84,32 +84,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
         },
     )
 
-    # assert_equals(solver.get_constraint_ids(), ["eq_capacity"])
+    # Add a brand new constraint
-    # assert_equals(
-    #     solver.get_constraint_rhs("eq_capacity"),
-    #     67.0,
-    # )
-    # assert_equals(
-    #     solver.get_constraint_lhs("eq_capacity"),
-    #     {
-    #         "x[0]": 23.0,
-    #         "x[1]": 26.0,
-    #         "x[2]": 20.0,
-    #         "x[3]": 18.0,
-    #     },
-    # )
-    # assert_equals(solver.get_constraint_sense("eq_capacity"), "<")
-
-    # if isinstance(solver, BasePyomoSolver):
-    #     model.cut = pe.Constraint(expr=model.x[0] <= 0.0, name="cut")
-    #     solver.add_constraint(model.cut)
-    # elif isinstance(solver, GurobiSolver):
-    #     x = model.getVarByName("x[0]")
-    #     solver.add_constraint(x <= 0.0, name="cut")
-    # else:
-    #     raise Exception("Illegal state")
-
-    # # Add a brand new constraint
     cut = instance.build_lazy_constraint(model, "cut")
     assert cut is not None
     solver.add_constraint(cut, name="cut")
@@ -140,44 +115,36 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # New constraint should affect the solution
     stats = solver.solve()
-    assert stats["Lower bound"] == 1030.0
+    assert_equals(stats["Lower bound"], 1030.0)
 
     # Verify slacks
-    assert solver.get_inequality_slacks() == {
+    assert_equals(
-        "cut": 0.0,
+        solver.get_inequality_slacks(),
-        "eq_capacity": 3.0,
+        {
-    }
+            "cut": 0.0,
+            "eq_capacity": 3.0,
+        },
+    )
 
     # # Extract the new constraint
-    # cobj = solver.extract_constraint("cut")
+    cobj = solver.extract_constraint("cut")
-    #
+
-    # # New constraint should no longer affect solution and should no longer
+    # New constraint should no longer affect solution
-    # # be listed in constraint ids
+    stats = solver.solve()
-    # assert solver.get_constraint_ids() == ["eq_capacity"]
+    assert_equals(stats["Lower bound"], 1183.0)
-    # stats = solver.solve()
+
-    # assert stats["Lower bound"] == 1183.0
+    # New constraint should not be satisfied by current solution
-    #
+    assert not solver.is_constraint_satisfied(cobj)
-    # # New constraint should not be satisfied by current solution
+
-    # assert not solver.is_constraint_satisfied(cobj)
+    # Re-add constraint
-    #
+    solver.add_constraint(cobj)
-    # # Re-add constraint
+
-    # solver.add_constraint(cobj)
+    # Constraint should affect solution again
-    #
+    stats = solver.solve()
-    # # Constraint should affect solution again
+    assert_equals(stats["Lower bound"], 1030.0)
-    # assert solver.get_constraint_ids() == ["eq_capacity", "cut"]
+
-    # stats = solver.solve()
+    # New constraint should now be satisfied
-    # assert stats["Lower bound"] == 1030.0
+    assert solver.is_constraint_satisfied(cobj)
-    #
-    # # New constraint should now be satisfied
-    # assert solver.is_constraint_satisfied(cobj)
-    #
-    # # Relax problem and make cut into an equality constraint
-    # solver.relax()
-    # solver.set_constraint_sense("cut", "=")
-    # stats = solver.solve()
-    # assert stats["Lower bound"] is not None
-    # assert round(stats["Lower bound"]) == 1030.0
-    # assert round(solver.get_dual("eq_capacity")) == 0.0
 
 
 def run_warm_start_tests(solver: InternalSolver) -> None:
@@ -195,8 +162,8 @@ def run_warm_start_tests(solver: InternalSolver) -> None:
 
     solver.fix({"x[0]": 1.0, "x[1]": 0.0, "x[2]": 0.0, "x[3]": 1.0})
     stats = solver.solve(tee=True)
-    assert stats["Lower bound"] == 725.0
+    assert_equals(stats["Lower bound"], 725.0)
-    assert stats["Upper bound"] == 725.0
+    assert_equals(stats["Upper bound"], 725.0)
 
 
 def run_infeasibility_tests(solver: InternalSolver) -> None:
@@ -223,7 +190,7 @@ def run_iteration_cb_tests(solver: InternalSolver) -> None:
         return count < 5
 
     solver.solve(iteration_cb=custom_iteration_cb)
-    assert count == 5
+    assert_equals(count, 5)
 
 
 def assert_equals(left: Any, right: Any) -> None: