Add types to internal solvers

miplearn/instance.py

@@ -164,12 +164,3 @@ class Instance(ABC):
         data = json.dumps(self.__dict__, indent=2).encode("utf-8")
         with gzip.GzipFile(filename, "w") as f:
             f.write(data)
-
-
-class PyomoInstance(Instance, ABC):
-    @abstractmethod
-    def to_model(self) -> pe.ConcreteModel:
-        """
-        Returns the concrete Pyomo model corresponding to this instance.
-        """
-        pass

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

miplearn/solvers/internal.py

@@ -13,6 +13,8 @@ from miplearn.types import (
     LazyCallback,
     MIPSolveStats,
     VarIndex,
+    Solution,
+    BranchPriorities,
 )
 
 logger = logging.getLogger(__name__)
@@ -79,7 +81,7 @@ class InternalSolver(ABC):
         pass
 
     @abstractmethod
-    def get_solution(self) -> Optional[Dict]:
+    def get_solution(self) -> Optional[Solution]:
         """
         Returns current solution found by the solver.
 
@@ -93,7 +95,7 @@ class InternalSolver(ABC):
         pass
 
     @abstractmethod
-    def set_warm_start(self, solution: Dict) -> None:
+    def set_warm_start(self, solution: Solution) -> None:
         """
         Sets the warm start to be used by the solver.
 
@@ -125,7 +127,7 @@ class InternalSolver(ABC):
         pass
 
     @abstractmethod
-    def fix(self, solution: Dict) -> None:
+    def fix(self, solution: Solution) -> None:
         """
         Fixes the values of a subset of decision variables.
 
@@ -135,7 +137,7 @@ class InternalSolver(ABC):
         """
         pass
 
-    def set_branching_priorities(self, priorities: Dict) -> None:
+    def set_branching_priorities(self, priorities: BranchPriorities) -> None:
         """
         Sets the branching priorities for the given decision variables.
 
@@ -147,7 +149,7 @@ class InternalSolver(ABC):
         `get_solution`. Missing values indicate variables whose priorities
         should not be modified.
         """
-        raise NotImplementedError()
+        raise Exception("Not implemented")
 
     @abstractmethod
     def get_constraint_ids(self) -> List[str]:

miplearn/solvers/pyomo/base.py

@@ -12,6 +12,7 @@ import pyomo
 from pyomo import environ as pe
 from pyomo.core import Var, Constraint
 from pyomo.opt import TerminationCondition
+from pyomo.opt.base.solvers import SolverFactory
 
 from miplearn.instance import Instance
 from miplearn.solvers import RedirectOutput
@@ -22,7 +23,7 @@ from miplearn.solvers.internal import (
     LazyCallback,
     MIPSolveStats,
 )
-from miplearn.types import VarIndex
+from miplearn.types import VarIndex, SolverParams, Solution
 
 logger = logging.getLogger(__name__)
 
@@ -34,19 +35,20 @@ class BasePyomoSolver(InternalSolver):
 
     def __init__(
         self,
-        solver_factory,
-        params,
-    ):
-        self.instance = None
-        self.model = None
-        self._all_vars = None
-        self._bin_vars = None
-        self._is_warm_start_available = False
-        self._pyomo_solver = solver_factory
-        self._obj_sense = None
-        self._varname_to_var = {}
-        self._cname_to_constr = {}
-        self._termination_condition = None
+        solver_factory: SolverFactory,
+        params: SolverParams,
+    ) -> None:
+        self.instance: Optional[Instance] = None
+        self.model: Optional[pe.ConcreteModel] = None
+        self._all_vars: List[pe.Var] = []
+        self._bin_vars: List[pe.Var] = []
+        self._is_warm_start_available: bool = False
+        self._pyomo_solver: SolverFactory = solver_factory
+        self._obj_sense: str = "min"
+        self._varname_to_var: Dict[str, pe.Var] = {}
+        self._cname_to_constr: Dict[str, pe.Constraint] = {}
+        self._termination_condition: str = ""
+
         for (key, value) in params.items():
             self._pyomo_solver.options[key] = value
 
@@ -88,8 +90,6 @@ class BasePyomoSolver(InternalSolver):
             streams += [sys.stdout]
         if iteration_cb is None:
             iteration_cb = lambda: False
-        self.instance.found_violated_lazy_constraints = []
-        self.instance.found_violated_user_cuts = []
         while True:
             logger.debug("Solving MIP...")
             with RedirectOutput(streams):
@@ -121,10 +121,11 @@ class BasePyomoSolver(InternalSolver):
         }
         return stats
 
-    def get_solution(self) -> Optional[Dict]:
+    def get_solution(self) -> Optional[Solution]:
+        assert self.model is not None
         if self.is_infeasible():
             return None
-        solution: Dict = {}
+        solution: Solution = {}
         for var in self.model.component_objects(Var):
             solution[str(var)] = {}
             for index in var:
@@ -133,7 +134,7 @@ class BasePyomoSolver(InternalSolver):
                 solution[str(var)][index] = var[index].value
         return solution
 
-    def set_warm_start(self, solution: Dict) -> None:
+    def set_warm_start(self, solution: Solution) -> None:
         self._clear_warm_start()
         count_total, count_fixed = 0, 0
         for var_name in solution:
@@ -172,8 +173,9 @@ class BasePyomoSolver(InternalSolver):
             var = self._varname_to_var[var_name]
             return var[index].value
 
-    def get_empty_solution(self) -> Dict:
-        solution: Dict = {}
+    def get_empty_solution(self) -> Solution:
+        assert self.model is not None
+        solution: Solution = {}
         for var in self.model.component_objects(Var):
             svar = str(var)
             solution[svar] = {}
@@ -195,6 +197,7 @@ class BasePyomoSolver(InternalSolver):
             self._obj_sense = "min"
 
     def _update_vars(self) -> None:
+        assert self.model is not None
         self._all_vars = []
         self._bin_vars = []
         self._varname_to_var = {}
@@ -206,6 +209,7 @@ class BasePyomoSolver(InternalSolver):
                     self._bin_vars += [var[idx]]
 
     def _update_constrs(self) -> None:
+        assert self.model is not None
         self._cname_to_constr = {}
         for constr in self.model.component_objects(Constraint):
             self._cname_to_constr[constr.name] = constr

miplearn/solvers/pyomo/cplex.py

@@ -1,11 +1,13 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
+from typing import Optional
 
 from pyomo import environ as pe
 from scipy.stats import randint
 
 from miplearn.solvers.pyomo.base import BasePyomoSolver
+from miplearn.types import SolverParams
 
 
 class CplexPyomoSolver(BasePyomoSolver):
@@ -19,13 +21,19 @@ class CplexPyomoSolver(BasePyomoSolver):
         {"mip_display": 5} to increase the log verbosity.
     """
 
-    def __init__(self, params=None):
+    def __init__(
+        self,
+        params: Optional[SolverParams] = None,
+    ) -> None:
+        if params is None:
+            params = {}
+        if "randomseed" not in params.keys():
+            params["randomseed"] = randint(low=0, high=1000).rvs()
+        if "mip_display" not in params.keys():
+            params["mip_display"] = 4
         super().__init__(
             solver_factory=pe.SolverFactory("cplex_persistent"),
-            params={
-                "randomseed": randint(low=0, high=1000).rvs(),
-                "mip_display": 4,
-            },
+            params=params,
         )
 
     def _get_warm_start_regexp(self):

miplearn/solvers/pyomo/gurobi.py

@@ -3,11 +3,13 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import logging
+from typing import Optional
 
 from pyomo import environ as pe
 from scipy.stats import randint
 
 from miplearn.solvers.pyomo.base import BasePyomoSolver
+from miplearn.types import SolverParams, BranchPriorities
 
 logger = logging.getLogger(__name__)
 
@@ -23,28 +25,35 @@ class GurobiPyomoSolver(BasePyomoSolver):
         {"Threads": 4} to set the number of threads.
     """
 
-    def __init__(self, params=None):
+    def __init__(
+        self,
+        params: SolverParams = None,
+    ) -> None:
+        if params is None:
+            params = {}
+        if "seed" not in params.keys():
+            params["seed"] = randint(low=0, high=1000).rvs()
         super().__init__(
             solver_factory=pe.SolverFactory("gurobi_persistent"),
-            params={
-                "Seed": randint(low=0, high=1000).rvs(),
-            },
+            params=params,
         )
 
-    def _extract_node_count(self, log):
+    def _extract_node_count(self, log: str) -> int:
         return max(1, int(self._pyomo_solver._solver_model.getAttr("NodeCount")))
 
-    def _get_warm_start_regexp(self):
+    def _get_warm_start_regexp(self) -> str:
         return "MIP start with objective ([0-9.e+-]*)"
 
-    def _get_node_count_regexp(self):
+    def _get_node_count_regexp(self) -> Optional[str]:
         return None
 
-    def set_branching_priorities(self, priorities):
+    def set_branching_priorities(self, priorities: BranchPriorities) -> None:
         from gurobipy import GRB
 
         for varname in priorities.keys():
             var = self._varname_to_var[varname]
             for (index, priority) in priorities[varname].items():
+                if priority is None:
+                    continue
                 gvar = self._pyomo_solver._pyomo_var_to_solver_var_map[var[index]]
                 gvar.setAttr(GRB.Attr.BranchPriority, int(round(priority)))

miplearn/solvers/pyomo/xpress.py

@@ -8,6 +8,7 @@ from pyomo import environ as pe
 from scipy.stats import randint
 
 from miplearn.solvers.pyomo.base import BasePyomoSolver
+from miplearn.types import SolverParams
 
 logger = logging.getLogger(__name__)
 
@@ -23,10 +24,12 @@ class XpressPyomoSolver(BasePyomoSolver):
         {"Threads": 4} to set the number of threads.
     """
 
-    def __init__(self, params=None):
+    def __init__(self, params: SolverParams = None) -> None:
+        if params is None:
+            params = {}
+        if "randomseed" not in params.keys():
+            params["randomseed"] = randint(low=0, high=1000).rvs()
         super().__init__(
             solver_factory=pe.SolverFactory("xpress_persistent"),
-            params={
-                "randomseed": randint(low=0, high=1000).rvs(),
-            },
+            params=params,
         )

miplearn/solvers/tests/__init__.py

@@ -7,7 +7,7 @@ from typing import List, Callable, Any
 
 from pyomo import environ as pe
 
-from miplearn.instance import Instance, PyomoInstance
+from miplearn.instance import Instance
 from miplearn.problems.knapsack import KnapsackInstance, GurobiKnapsackInstance
 from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.internal import InternalSolver
@@ -16,7 +16,7 @@ from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
 from miplearn.solvers.pyomo.xpress import XpressPyomoSolver
 
 
-class InfeasiblePyomoInstance(PyomoInstance):
+class InfeasiblePyomoInstance(Instance):
     def to_model(self) -> pe.ConcreteModel:
         model = pe.ConcreteModel()
         model.x = pe.Var([0], domain=pe.Binary)

miplearn/types.py

@@ -6,15 +6,19 @@ from typing import Optional, Dict, Callable, Any, Union, List
 
 from mypy_extensions import TypedDict
 
+VarIndex = Union[str, int, List[Union[str, int]]]
+
+Solution = Dict[str, Dict[VarIndex, Optional[float]]]
+
 TrainingSample = TypedDict(
     "TrainingSample",
     {
         "LP log": str,
-        "LP solution": Optional[Dict],
+        "LP solution": Optional[Solution],
         "LP value": Optional[float],
         "Lower bound": Optional[float],
         "MIP log": str,
-        "Solution": Optional[Dict],
+        "Solution": Optional[Solution],
         "Upper bound": Optional[float],
         "slacks": Dict,
     },
@@ -47,4 +51,6 @@ IterationCallback = Callable[[], bool]
 
 LazyCallback = Callable[[Any, Any], None]
 
-VarIndex = Union[str, int, List[Union[str, int]]]
+SolverParams = Dict[str, Any]
+
+BranchPriorities = Solution