Lazy: Simplify method signature; switch to AbstractModel

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
+from typing import Optional, Dict, Callable
 
 import numpy as np
 
@@ -16,6 +16,10 @@ class AbstractModel(ABC):
     _supports_node_count = False
     _supports_solution_pool = False
 
+    def __init__(self) -> None:
+        self.lazy_enforce: Optional[Callable] = None
+        self.lazy_separate: Optional[Callable] = None
+
     @abstractmethod
     def add_constrs(
         self,

miplearn/solvers/gurobi.py

@@ -28,6 +28,7 @@ class GurobiModel(AbstractModel):
         self.lazy_enforce = lazy_enforce
         self.inner = inner
         self.lazy_constrs_: Optional[List[Any]] = None
+        self.where = "default"
 
     def add_constrs(
         self,
@@ -53,6 +54,14 @@ class GurobiModel(AbstractModel):
                 stats["Added constraints"] = 0
             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)
+
     def extract_after_load(self, h5: H5File) -> None:
         """
         Given a model that has just been loaded, extracts static problem
@@ -132,9 +141,11 @@ class GurobiModel(AbstractModel):
             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, "cb")
+                self.lazy_enforce(self, violations)
+            self.where = "default"
 
         if self.lazy_enforce is not None:
             self.inner.Params.lazyConstraints = 1