From 5e1f26e4b00113fc55a1018c5803e900494a80c1 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Sat, 10 Apr 2021 17:38:03 -0500
Subject: [PATCH] Add more constraint features

---
 miplearn/solvers/gurobi.py         | 26 ++++++++++++++++--
 miplearn/solvers/internal.py       |  7 +++++
 miplearn/solvers/pyomo/base.py     | 13 +++++++++
 miplearn/solvers/tests/__init__.py | 44 ++++++++++++++++++++++--------
 4 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/miplearn/solvers/gurobi.py b/miplearn/solvers/gurobi.py
index da83837..500f26b 100644
--- a/miplearn/solvers/gurobi.py
+++ b/miplearn/solvers/gurobi.py
@@ -4,8 +4,6 @@
 import logging
 import re
 import sys
-from dataclasses import dataclass
-from enum import Enum
 from io import StringIO
 from random import randint
 from typing import List, Any, Dict, Optional, Hashable
@@ -444,6 +442,14 @@ class GurobiSolver(InternalSolver):
         )
         if self._has_lp_solution:
             constr.dual_value = gp_constr.pi
+            constr.sa_rhs_up = gp_constr.sarhsup
+            constr.sa_rhs_low = gp_constr.sarhslow
+            if gp_constr.cbasis == 0:
+                constr.basis_status = "b"
+            elif gp_constr.cbasis == -1:
+                constr.basis_status = "n"
+            else:
+                raise Exception(f"unknown cbasis: {gp_constr.cbasis}")
         if self._has_lp_solution or self._has_mip_solution:
             constr.slack = gp_constr.slack
         return constr
@@ -452,6 +458,22 @@ class GurobiSolver(InternalSolver):
     def are_callbacks_supported(self) -> bool:
         return True
 
+    @overrides
+    def get_constraint_attrs(self) -> List[str]:
+        return [
+            "basis_status",
+            "category",
+            "dual_value",
+            "lazy",
+            "lhs",
+            "rhs",
+            "sa_rhs_down",
+            "sa_rhs_up",
+            "sense",
+            "slack",
+            "user_features",
+        ]
+
 
 class GurobiTestInstanceInfeasible(Instance):
     @overrides
diff --git a/miplearn/solvers/internal.py b/miplearn/solvers/internal.py
index 0149549..c4589b2 100644
--- a/miplearn/solvers/internal.py
+++ b/miplearn/solvers/internal.py
@@ -246,3 +246,10 @@ class InternalSolver(ABC, EnforceOverrides):
         callback or user cuts callback.
         """
         return False
+
+    @abstractmethod
+    def get_constraint_attrs(self) -> List[str]:
+        """
+        Returns a list of constraint attributes supported by this solver.
+        """
+        pass
diff --git a/miplearn/solvers/pyomo/base.py b/miplearn/solvers/pyomo/base.py
index 2785985..e472f67 100644
--- a/miplearn/solvers/pyomo/base.py
+++ b/miplearn/solvers/pyomo/base.py
@@ -435,6 +435,19 @@ class BasePyomoSolver(InternalSolver):
     def are_callbacks_supported(self) -> bool:
         return False
 
+    @overrides
+    def get_constraint_attrs(self) -> List[str]:
+        return [
+            "category",
+            "dual_value",
+            "lazy",
+            "lhs",
+            "rhs",
+            "sense",
+            "slack",
+            "user_features",
+        ]
+
 
 class PyomoTestInstanceInfeasible(Instance):
     @overrides
diff --git a/miplearn/solvers/tests/__init__.py b/miplearn/solvers/tests/__init__.py
index 2637075..083de11 100644
--- a/miplearn/solvers/tests/__init__.py
+++ b/miplearn/solvers/tests/__init__.py
@@ -2,7 +2,7 @@
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from typing import Any
+from typing import Any, Dict
 
 from miplearn.features import Constraint
 from miplearn.solvers.internal import InternalSolver
@@ -12,7 +12,7 @@ from miplearn.solvers.internal import InternalSolver
 # This file is in the main source folder, so that it can be called from Julia.
 
 
-def _round_constraints(constraints):
+def _round_constraints(constraints: Dict[str, Constraint]) -> Dict[str, Constraint]:
     for (cname, c) in constraints.items():
         for attr in ["slack", "dual_value"]:
             if getattr(c, attr) is not None:
@@ -20,6 +20,20 @@ def _round_constraints(constraints):
     return constraints
 
 
+def _remove_unsupported_constr_attrs(
+    solver: InternalSolver,
+    constraints: Dict[str, Constraint],
+):
+    for (cname, c) in constraints.items():
+        to_remove = []
+        for k in c.__dict__.keys():
+            if k not in solver.get_constraint_attrs():
+                to_remove.append(k)
+        for k in to_remove:
+            setattr(c, k, None)
+    return constraints
+
+
 def run_internal_solver_tests(solver: InternalSolver) -> None:
     run_basic_usage_tests(solver.clone())
     run_warm_start_tests(solver.clone())
@@ -76,16 +90,22 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
     # Fetch constraints (after-lp)
     assert_equals(
         _round_constraints(solver.get_constraints()),
-        {
-            "eq_capacity": Constraint(
-                lazy=False,
-                lhs={"x[0]": 23.0, "x[1]": 26.0, "x[2]": 20.0, "x[3]": 18.0},
-                rhs=67.0,
-                sense="<",
-                slack=0.0,
-                dual_value=13.538462,
-            )
-        },
+        _remove_unsupported_constr_attrs(
+            solver,
+            {
+                "eq_capacity": Constraint(
+                    lazy=False,
+                    lhs={"x[0]": 23.0, "x[1]": 26.0, "x[2]": 20.0, "x[3]": 18.0},
+                    rhs=67.0,
+                    sense="<",
+                    slack=0.0,
+                    dual_value=13.538462,
+                    sa_rhs_down=None,
+                    sa_rhs_up=69.0,
+                    basis_status="n",
+                )
+            },
+        ),
     )
 
     # Solve MIP