Lazy: Minor fixes; make it compatible with Pyomo

miplearn/solvers/abstract.py

@@ -3,7 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 from abc import ABC, abstractmethod
-from typing import Optional, Dict, Callable
+from typing import Optional, Dict, Callable, Hashable, List, Any
 
 import numpy as np
 
@@ -16,9 +16,15 @@ class AbstractModel(ABC):
     _supports_node_count = False
     _supports_solution_pool = False
 
+    WHERE_DEFAULT = "default"
+    WHERE_CUTS = "cuts"
+    WHERE_LAZY = "lazy"
+
     def __init__(self) -> None:
         self.lazy_enforce: Optional[Callable] = None
         self.lazy_separate: Optional[Callable] = None
+        self.lazy_constrs_: Optional[List[Any]] = None
+        self.where = self.WHERE_DEFAULT
 
     @abstractmethod
     def add_constrs(

miplearn/solvers/gurobi.py

@@ -12,6 +12,27 @@ from miplearn.h5 import H5File
 from miplearn.solvers.abstract import AbstractModel
 
 
+def _gurobi_callback(model: AbstractModel, where: int) -> None:
+    assert model.lazy_separate is not None
+    assert model.lazy_enforce is not None
+    assert model.lazy_constrs_ is not None
+    if where == GRB.Callback.MIPSOL:
+        model.where = model.WHERE_LAZY
+        violations = model.lazy_separate(model)
+        model.lazy_constrs_.extend(violations)
+        model.lazy_enforce(model, violations)
+    model.where = model.WHERE_DEFAULT
+
+
+def _gurobi_add_constr(gp_model: gp.Model, where: str, constr: Any) -> None:
+    if where == AbstractModel.WHERE_LAZY:
+        gp_model.cbLazy(constr)
+    elif where == AbstractModel.WHERE_CUTS:
+        gp_model.cbCut(constr)
+    else:
+        gp_model.addConstr(constr)
+
+
 class GurobiModel(AbstractModel):
     _supports_basis_status = True
     _supports_sensitivity_analysis = True
@@ -24,11 +45,10 @@ class GurobiModel(AbstractModel):
         lazy_separate: Optional[Callable] = None,
         lazy_enforce: Optional[Callable] = None,
     ) -> None:
+        super().__init__()
         self.lazy_separate = lazy_separate
         self.lazy_enforce = lazy_enforce
         self.inner = inner
-        self.lazy_constrs_: Optional[List[Any]] = None
-        self.where = "default"
 
     def add_constrs(
         self,
@@ -55,12 +75,7 @@ class GurobiModel(AbstractModel):
             stats["Added constraints"] += nconstrs
 
     def add_constr(self, constr: Any) -> None:
-        if self.where == "lazy":
-            self.inner.cbLazy(constr)
-        elif self.where == "cut":
-            self.inner.cbCut(constr)
-        else:
-            self.inner.addConstr(constr)
+        _gurobi_add_constr(self.inner, self.where, constr)
 
     def extract_after_load(self, h5: H5File) -> None:
         """
@@ -136,19 +151,12 @@ class GurobiModel(AbstractModel):
     def optimize(self) -> None:
         self.lazy_constrs_ = []
 
-        def callback(m: gp.Model, where: int) -> None:
-            assert self.lazy_separate is not None
-            assert self.lazy_constrs_ is not None
-            assert self.lazy_enforce is not None
-            if where == GRB.Callback.MIPSOL:
-                self.where = "lazy"
-                violations = self.lazy_separate(self)
-                self.lazy_constrs_.extend(violations)
-                self.lazy_enforce(self, violations)
-            self.where = "default"
+        def callback(_: gp.Model, where: int) -> None:
+            _gurobi_callback(self, where)
 
         if self.lazy_enforce is not None:
-            self.inner.Params.lazyConstraints = 1
+            self.inner.setParam("PreCrush", 1)
+            self.inner.setParam("LazyConstraints", 1)
             self.inner.optimize(callback)
         else:
             self.inner.optimize()

miplearn/solvers/pyomo.py

@@ -2,10 +2,11 @@
 #  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 from numbers import Number
-from typing import Optional, Dict, List, Any, Tuple, Union
+from typing import Optional, Dict, List, Any, Tuple, Callable
 
 import numpy as np
 import pyomo
+import pyomo.environ as pe
 from pyomo.core import Objective, Var, Suffix
 from pyomo.core.base import _GeneralVarData
 from pyomo.core.expr.numeric_expr import SumExpression, MonomialTermExpression
@@ -13,11 +14,18 @@ from scipy.sparse import coo_matrix
 
 from miplearn.h5 import H5File
 from miplearn.solvers.abstract import AbstractModel
-import pyomo.environ as pe
+from miplearn.solvers.gurobi import _gurobi_callback, _gurobi_add_constr
 
 
 class PyomoModel(AbstractModel):
-    def __init__(self, model: pe.ConcreteModel, solver_name: str = "gurobi_persistent"):
+    def __init__(
+        self,
+        model: pe.ConcreteModel,
+        solver_name: str = "gurobi_persistent",
+        lazy_separate: Optional[Callable] = None,
+        lazy_enforce: Optional[Callable] = None,
+    ):
+        super().__init__()
         self.inner = model
         self.solver_name = solver_name
         self.solver = pe.SolverFactory(solver_name)
@@ -26,11 +34,20 @@ class PyomoModel(AbstractModel):
             self.solver.set_instance(model)
         self.results: Optional[Dict] = None
         self._is_warm_start_available = False
+        self.lazy_separate = lazy_separate
+        self.lazy_enforce = lazy_enforce
+        self.lazy_constrs_: Optional[List[Any]] = None
         if not hasattr(self.inner, "dual"):
             self.inner.dual = Suffix(direction=Suffix.IMPORT)
             self.inner.rc = Suffix(direction=Suffix.IMPORT)
             self.inner.slack = Suffix(direction=Suffix.IMPORT)
 
+    def add_constr(self, constr: Any) -> None:
+        assert (
+            self.solver_name == "gurobi_persistent"
+        ), "Callbacks are currently only supported on gurobi_persistent"
+        _gurobi_add_constr(self.solver, self.where, constr)
+
     def add_constrs(
         self,
         var_names: np.ndarray,
@@ -114,6 +131,20 @@ class PyomoModel(AbstractModel):
                 self.solver.update_var(var)
 
     def optimize(self) -> None:
+        self.lazy_constrs_ = []
+
+        if self.lazy_separate is not None:
+            assert (
+                self.solver_name == "gurobi_persistent"
+            ), "Callbacks are currently only supported on gurobi_persistent"
+
+            def callback(_: Any, __: Any, where: int) -> None:
+                _gurobi_callback(self, where)
+
+            self.solver.set_gurobi_param("PreCrush", 1)
+            self.solver.set_gurobi_param("LazyConstraints", 1)
+            self.solver.set_callback(callback)
+
         if self.is_persistent:
             self.results = self.solver.solve(
                 tee=True,