Lazy: Minor fixes; make it compatible with Pyomo

miplearn/collectors/basic.py

@@ -4,6 +4,8 @@
 
 import json
 import os
+import sys
+
 from io import StringIO
 from os.path import exists
 from typing import Callable, List
@@ -57,11 +59,8 @@ class BasicCollector:
                     model.extract_after_mip(h5)
 
                     # Add lazy constraints to model
-                    if (
-                        hasattr(model, "lazy_enforce")
-                        and model.lazy_enforce is not None
-                    ):
-                        model.lazy_enforce(model, model.lazy_constrs_, "aot")
+                    if model.lazy_enforce is not None:
+                        model.lazy_enforce(model, model.lazy_constrs_)
                         h5.put_scalar("mip_lazy", repr(model.lazy_constrs_))
 
                     # Save MPS file

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,

tests/test_lazy_pyomo.py

@@ -0,0 +1,44 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2023, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+import logging
+from typing import Any, Hashable, List
+
+import pyomo.environ as pe
+
+from miplearn.solvers.pyomo import PyomoModel
+
+logger = logging.getLogger(__name__)
+
+
+def _build_model() -> PyomoModel:
+    m = pe.ConcreteModel()
+    m.x = pe.Var(bounds=(0, 5), domain=pe.Integers)
+    m.obj = pe.Objective(expr=-m.x)
+    m.cons = pe.ConstraintList()
+
+    def lazy_separate(model: PyomoModel) -> List[Hashable]:
+        model.solver.cbGetSolution(vars=[m.x])
+        if m.x.value > 0.5:
+            return [m.x.value]
+        else:
+            return []
+
+    def lazy_enforce(model: PyomoModel, violations: List[Any]) -> None:
+        for v in violations:
+            model.add_constr(m.cons.add(m.x <= round(v - 1)))
+
+    return PyomoModel(
+        m,
+        "gurobi_persistent",
+        lazy_separate=lazy_separate,
+        lazy_enforce=lazy_enforce,
+    )
+
+
+def test_pyomo_callback() -> None:
+    model = _build_model()
+    model.optimize()
+    assert model.lazy_constrs_ is not None
+    assert len(model.lazy_constrs_) > 0
+    assert model.inner.x.value == 0.0