From 868675ecf237a083deb238d8e84261a7b8442e70 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Thu, 21 Jan 2021 14:24:06 -0600
Subject: [PATCH] Implement some constraint methods in Pyomo

---
 miplearn/solvers/internal.py                  |  4 --
 miplearn/solvers/pyomo/base.py                | 38 +++++++++++++------
 .../solvers/tests/test_internal_solver.py     | 19 +++++-----
 3 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/miplearn/solvers/internal.py b/miplearn/solvers/internal.py
index 792ac5e..30e0092 100644
--- a/miplearn/solvers/internal.py
+++ b/miplearn/solvers/internal.py
@@ -18,10 +18,6 @@ from miplearn.types import (
 logger = logging.getLogger(__name__)
 
 
-class ExtractedConstraint(ABC):
-    pass
-
-
 class Constraint:
     pass
 
diff --git a/miplearn/solvers/pyomo/base.py b/miplearn/solvers/pyomo/base.py
index defea3d..25c34fb 100644
--- a/miplearn/solvers/pyomo/base.py
+++ b/miplearn/solvers/pyomo/base.py
@@ -270,12 +270,6 @@ class BasePyomoSolver(InternalSolver):
     def _get_node_count_regexp(self) -> Optional[str]:
         return None
 
-    def extract_constraint(self, cid):
-        raise Exception("Not implemented")
-
-    def is_constraint_satisfied(self, cobj):
-        raise Exception("Not implemented")
-
     def relax(self) -> None:
         for var in self._bin_vars:
             lb, ub = var.bounds
@@ -285,15 +279,35 @@ class BasePyomoSolver(InternalSolver):
             self._pyomo_solver.update_var(var)
 
     def get_inequality_slacks(self) -> Dict[str, float]:
-        raise Exception("not implemented")
+        result: Dict[str, float] = {}
+        for (cname, cobj) in self._cname_to_constr.items():
+            if cobj.equality:
+                continue
+            result[cname] = cobj.slack()
+        return result
+
+    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), "range constraints not supported"
+        if has_lb:
+            return ">"
+        elif has_ub:
+            return "<"
+        else:
+            return "="
+
+    def set_constraint_sense(self, cid: str, sense: str) -> None:
+        raise Exception("Not implemented")
 
-    def set_constraint_sense(self, cid, sense):
+    def extract_constraint(self, cid: str) -> Constraint:
         raise Exception("Not implemented")
 
-    def get_constraint_sense(self, cid):
+    def is_constraint_satisfied(self, cobj: Constraint) -> bool:
         raise Exception("Not implemented")
 
-    def set_constraint_rhs(self, cid, rhs):
+    def set_constraint_rhs(self, cid: str, rhs: float) -> None:
         raise Exception("Not implemented")
 
     def is_infeasible(self) -> bool:
@@ -302,5 +316,5 @@ class BasePyomoSolver(InternalSolver):
     def get_dual(self, cid):
         raise Exception("Not implemented")
 
-    def get_sense(self):
-        raise Exception("Not implemented")
+    def get_sense(self) -> str:
+        return self._obj_sense
diff --git a/miplearn/solvers/tests/test_internal_solver.py b/miplearn/solvers/tests/test_internal_solver.py
index 4356e4e..98c763b 100644
--- a/miplearn/solvers/tests/test_internal_solver.py
+++ b/miplearn/solvers/tests/test_internal_solver.py
@@ -133,18 +133,17 @@ def test_internal_solver():
         stats = solver.solve()
         assert stats["Lower bound"] == 1030.0
 
-        if isinstance(solver, GurobiSolver):
-
-            assert solver.get_sense() == "max"
-            assert solver.get_constraint_sense("cut") == "<"
-            assert solver.get_constraint_sense("eq_capacity") == "<"
+        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,
-                "eq_capacity": 3.0,
-            }
+        # Verify slacks
+        assert solver.get_inequality_slacks() == {
+            "cut": 0.0,
+            "eq_capacity": 3.0,
+        }
 
+        if isinstance(solver, GurobiSolver):
             # Extract the new constraint
             cobj = solver.extract_constraint("cut")