Add types to internal solvers

miplearn/solvers/gurobi.py

@@ -6,7 +6,7 @@ import re
 import sys
 from io import StringIO
 from random import randint
-from typing import List, Any, Dict, Union, Tuple, Optional
+from typing import List, Any, Dict, Optional
 
 from miplearn.instance import Instance
 from miplearn.solvers import RedirectOutput
@@ -17,7 +17,7 @@ from miplearn.solvers.internal import (
     LazyCallback,
     MIPSolveStats,
 )
-from miplearn.types import VarIndex
+from miplearn.types import VarIndex, SolverParams, Solution
 
 logger = logging.getLogger(__name__)
 
@@ -25,9 +25,9 @@ logger = logging.getLogger(__name__)
 class GurobiSolver(InternalSolver):
     def __init__(
         self,
-        params=None,
-        lazy_cb_frequency=1,
-    ):
+        params: Optional[SolverParams] = None,
+        lazy_cb_frequency: int = 1,
+    ) -> None:
         """
         An InternalSolver backed by Gurobi's Python API (without Pyomo).
 
@@ -41,24 +41,28 @@ class GurobiSolver(InternalSolver):
             is found. If 2, calls it also at every node, after solving the
             LP relaxation of that node.
         """
+        import gurobipy
+
         if params is None:
             params = {}
         params["InfUnbdInfo"] = True
-        import gurobipy
 
         self.gp = gurobipy
-        self.GRB = gurobipy.GRB
-        self.instance = None
-        self.model = None
-        self.params = params
+        self.instance: Optional[Instance] = None
+        self.model: Optional["gurobipy.Model"] = None
+        self.params: SolverParams = params
         self._all_vars: Dict = {}
-        self._bin_vars = None
-        self.cb_where = None
+        self._bin_vars: Optional[Dict[str, Dict[VarIndex, "gurobipy.Var"]]] = None
+        self.cb_where: Optional[int] = None
+
         assert lazy_cb_frequency in [1, 2]
         if lazy_cb_frequency == 1:
-            self.lazy_cb_where = [self.GRB.Callback.MIPSOL]
+            self.lazy_cb_where = [self.gp.GRB.Callback.MIPSOL]
         else:
-            self.lazy_cb_where = [self.GRB.Callback.MIPSOL, self.GRB.Callback.MIPNODE]
+            self.lazy_cb_where = [
+                self.gp.GRB.Callback.MIPSOL,
+                self.gp.GRB.Callback.MIPNODE,
+            ]
 
     def set_instance(
         self,
@@ -79,9 +83,10 @@ class GurobiSolver(InternalSolver):
             raise Exception("method cannot be called from a callback")
 
     def _update_vars(self) -> None:
+        assert self.model is not None
         self._all_vars = {}
         self._bin_vars = {}
-        idx: Union[Tuple, List[int], int]
+        idx: VarIndex
         for var in self.model.getVars():
             m = re.search(r"([^[]*)\[(.*)]", var.varName)
             if m is None:
@@ -89,9 +94,8 @@ class GurobiSolver(InternalSolver):
                 idx = [0]
             else:
                 name = m.group(1)
-                idx = tuple(
-                    int(k) if k.isdecimal() else k for k in m.group(2).split(",")
-                )
+                parts = m.group(2).split(",")
+                idx = [int(k) if k.isdecimal else k for k in parts]
             if len(idx) == 1:
                 idx = idx[0]
             if name not in self._all_vars:
@@ -103,6 +107,7 @@ class GurobiSolver(InternalSolver):
                 self._bin_vars[name][idx] = var
 
     def _apply_params(self, streams: List[Any]) -> None:
+        assert self.model is not None
         with RedirectOutput(streams):
             for (name, value) in self.params.items():
                 self.model.setParam(name, value)
@@ -118,16 +123,18 @@ class GurobiSolver(InternalSolver):
         if tee:
             streams += [sys.stdout]
         self._apply_params(streams)
+        assert self.model is not None
+        assert self._bin_vars is not None
         for (varname, vardict) in self._bin_vars.items():
             for (idx, var) in vardict.items():
-                var.vtype = self.GRB.CONTINUOUS
+                var.vtype = self.gp.GRB.CONTINUOUS
                 var.lb = 0.0
                 var.ub = 1.0
         with RedirectOutput(streams):
             self.model.optimize()
         for (varname, vardict) in self._bin_vars.items():
             for (idx, var) in vardict.items():
-                var.vtype = self.GRB.BINARY
+                var.vtype = self.gp.GRB.BINARY
         log = streams[0].getvalue()
         opt_value = None
         if not self.is_infeasible():
@@ -144,6 +151,7 @@ class GurobiSolver(InternalSolver):
         lazy_cb: LazyCallback = None,
     ) -> MIPSolveStats:
         self._raise_if_callback()
+        assert self.model is not None
 
         def cb_wrapper(cb_model, cb_where):
             try:
@@ -199,18 +207,19 @@ class GurobiSolver(InternalSolver):
         }
         return stats
 
-    def get_solution(self) -> Optional[Dict]:
+    def get_solution(self) -> Optional[Solution]:
         self._raise_if_callback()
+        assert self.model is not None
         if self.model.solCount == 0:
             return None
-        solution: Dict = {}
+        solution: Solution = {}
         for (varname, vardict) in self._all_vars.items():
             solution[varname] = {}
             for (idx, var) in vardict.items():
                 solution[varname][idx] = var.x
         return solution
 
-    def set_warm_start(self, solution: Dict) -> None:
+    def set_warm_start(self, solution: Solution) -> None:
         self._raise_if_callback()
         self._clear_warm_start()
         count_fixed, count_total = 0, 0
@@ -225,20 +234,27 @@ class GurobiSolver(InternalSolver):
             % (count_fixed, count_total)
         )
 
-    def get_sense(self):
+    def get_sense(self) -> str:
+        assert self.model is not None
         if self.model.modelSense == 1:
             return "min"
         else:
             return "max"
 
-    def get_value(self, var_name: str, index: VarIndex) -> Optional[float]:
+    def get_value(
+        self,
+        var_name: str,
+        index: VarIndex,
+    ) -> Optional[float]:
         var = self._all_vars[var_name][index]
         return self._get_value(var)
 
     def is_infeasible(self) -> bool:
-        return self.model.status in [self.GRB.INFEASIBLE, self.GRB.INF_OR_UNBD]
+        assert self.model is not None
+        return self.model.status in [self.gp.GRB.INFEASIBLE, self.gp.GRB.INF_OR_UNBD]
 
-    def get_dual(self, cid):
+    def get_dual(self, cid: str) -> float:
+        assert self.model is not None
         c = self.model.getConstrByName(cid)
         if self.is_infeasible():
             return c.farkasDual
@@ -246,9 +262,10 @@ class GurobiSolver(InternalSolver):
             return c.pi
 
     def _get_value(self, var: Any) -> Optional[float]:
-        if self.cb_where == self.GRB.Callback.MIPSOL:
+        assert self.model is not None
+        if self.cb_where == self.gp.GRB.Callback.MIPSOL:
             return self.model.cbGetSolution(var)
-        elif self.cb_where == self.GRB.Callback.MIPNODE:
+        elif self.cb_where == self.gp.GRB.Callback.MIPNODE:
             return self.model.cbGetNodeRel(var)
         elif self.cb_where is None:
             if self.is_infeasible():
@@ -260,24 +277,35 @@ class GurobiSolver(InternalSolver):
                 "get_value cannot be called from cb_where=%s" % self.cb_where
             )
 
-    def get_empty_solution(self) -> Dict:
+    def get_empty_solution(self) -> Solution:
         self._raise_if_callback()
-        solution: Dict = {}
+        solution: Solution = {}
         for (varname, vardict) in self._all_vars.items():
             solution[varname] = {}
             for (idx, var) in vardict.items():
                 solution[varname][idx] = None
         return solution
 
-    def add_constraint(self, constraint, name=""):
+    def add_constraint(
+        self,
+        constraint: Any,
+        name: str = "",
+    ) -> None:
+        assert self.model is not None
         if type(constraint) is tuple:
             lhs, sense, rhs, name = constraint
-            if self.cb_where in [self.GRB.Callback.MIPSOL, self.GRB.Callback.MIPNODE]:
+            if self.cb_where in [
+                self.gp.GRB.Callback.MIPSOL,
+                self.gp.GRB.Callback.MIPNODE,
+            ]:
                 self.model.cbLazy(lhs, sense, rhs)
             else:
                 self.model.addConstr(lhs, sense, rhs, name)
         else:
-            if self.cb_where in [self.GRB.Callback.MIPSOL, self.GRB.Callback.MIPNODE]:
+            if self.cb_where in [
+                self.gp.GRB.Callback.MIPSOL,
+                self.gp.GRB.Callback.MIPNODE,
+            ]:
                 self.model.cbLazy(constraint)
             else:
                 self.model.addConstr(constraint, name=name)
@@ -285,16 +313,16 @@ class GurobiSolver(InternalSolver):
     def _clear_warm_start(self) -> None:
         for (varname, vardict) in self._all_vars.items():
             for (idx, var) in vardict.items():
-                var.start = self.GRB.UNDEFINED
+                var.start = self.gp.GRB.UNDEFINED
 
-    def fix(self, solution):
+    def fix(self, solution: Solution) -> None:
         self._raise_if_callback()
         for (varname, vardict) in solution.items():
             for (idx, value) in vardict.items():
                 if value is None:
                     continue
                 var = self._all_vars[varname][idx]
-                var.vtype = self.GRB.CONTINUOUS
+                var.vtype = self.gp.GRB.CONTINUOUS
                 var.lb = value
                 var.ub = value
 
@@ -330,6 +358,7 @@ class GurobiSolver(InternalSolver):
             raise Exception("Unknown sense: %s" % sense)
 
     def get_inequality_slacks(self) -> Dict[str, float]:
+        assert self.model is not None
         ineqs = [c for c in self.model.getConstrs() if c.sense != "="]
         return {c.ConstrName: c.Slack for c in ineqs}
 
@@ -342,6 +371,7 @@ class GurobiSolver(InternalSolver):
         return c.Sense
 
     def relax(self) -> None:
+        assert self.model is not None
         self.model = self.model.relax()
         self._update_vars()
 
@@ -372,11 +402,9 @@ class GurobiSolver(InternalSolver):
         }
 
     def __setstate__(self, state):
-        from gurobipy import GRB
 
         self.params = state["params"]
         self.lazy_cb_where = state["lazy_cb_where"]
-        self.GRB = GRB
         self.instance = None
         self.model = None
         self._all_vars = None