Make lazy constr component compatible with Pyomo+Gurobi

miplearn/problems/__init__.py

@@ -1,3 +1,28 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
+
+from typing import Any, Optional
+
+import gurobipy as gp
+from pyomo import environ as pe
+
+
+def _gurobipy_set_params(model: gp.Model, params: Optional[dict[str, Any]]) -> None:
+    assert isinstance(model, gp.Model)
+    if params is not None:
+        for (param_name, param_value) in params.items():
+            setattr(model.params, param_name, param_value)
+
+
+def _pyomo_set_params(
+    model: pe.ConcreteModel,
+    params: Optional[dict[str, Any]],
+    solver: str,
+) -> None:
+    assert (
+        solver == "gurobi_persistent"
+    ), "setting parameters is only supported with gurobi_persistent"
+    if solver == "gurobi_persistent" and params is not None:
+        for (param_name, param_value) in params.items():
+            model.solver.set_gurobi_param(param_name, param_value)

miplearn/problems/stab.py

@@ -18,6 +18,8 @@ from networkx import Graph
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
 
+from . import _gurobipy_set_params, _pyomo_set_params
+
 logger = logging.getLogger(__name__)
 
 
@@ -88,11 +90,10 @@ def build_stab_model_gurobipy(
     data: Union[str, MaxWeightStableSetData],
     params: Optional[dict[str, Any]] = None,
 ) -> GurobiModel:
-    data = _stab_read(data)
     model = gp.Model()
-    if params is not None:
-        for (param_name, param_value) in params.items():
-            setattr(model.params, param_name, param_value)
+    _gurobipy_set_params(model, params)
+
+    data = _stab_read(data)
     nodes = list(data.graph.nodes)
 
     # Variables and objective function
@@ -152,18 +153,14 @@ def build_stab_model_pyomo(
         for clique in violations:
             m.add_constr(model.clique_eqs.add(sum(model.x[i] for i in clique) <= 1))
 
-    m = PyomoModel(
+    pm = PyomoModel(
         model,
         solver,
         cuts_separate=cuts_separate,
         cuts_enforce=cuts_enforce,
     )
-
-    if solver == "gurobi_persistent" and params is not None:
-        for (param_name, param_value) in params.items():
-            m.solver.set_gurobi_param(param_name, param_value)
-
-    return m
+    _pyomo_set_params(pm, params, solver)
+    return pm
 
 
 def _stab_read(data: Union[str, MaxWeightStableSetData]) -> MaxWeightStableSetData:

miplearn/problems/tsp.py

@@ -2,20 +2,23 @@
 #  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
 from dataclasses import dataclass
 from typing import List, Tuple, Optional, Any, Union
 
 import gurobipy as gp
 import networkx as nx
 import numpy as np
+import pyomo.environ as pe
 from gurobipy import quicksum, GRB, tuplelist
+from miplearn.io import read_pkl_gz
+from miplearn.problems import _gurobipy_set_params, _pyomo_set_params
+from miplearn.solvers.gurobi import GurobiModel
 from scipy.spatial.distance import pdist, squareform
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
-import logging
 
-from miplearn.io import read_pkl_gz
-from miplearn.solvers.gurobi import GurobiModel
+from miplearn.solvers.pyomo import PyomoModel
 
 logger = logging.getLogger(__name__)
 
@@ -112,15 +115,18 @@ class TravelingSalesmanGenerator:
         return n, cities
 
 
-def build_tsp_model(data: Union[str, TravelingSalesmanData]) -> GurobiModel:
-    if isinstance(data, str):
-        data = read_pkl_gz(data)
-    assert isinstance(data, TravelingSalesmanData)
+def build_tsp_model_gurobipy(
+    data: Union[str, TravelingSalesmanData],
+    params: Optional[dict[str, Any]] = None,
+) -> GurobiModel:
 
+    model = gp.Model()
+    _gurobipy_set_params(model, params)
+
+    data = _tsp_read(data)
     edges = tuplelist(
         (i, j) for i in range(data.n_cities) for j in range(i + 1, data.n_cities)
     )
-    model = gp.Model()
 
     # Decision variables
     x = model.addVars(edges, vtype=GRB.BINARY, name="x")
@@ -173,3 +179,75 @@ def build_tsp_model(data: Union[str, TravelingSalesmanData]) -> GurobiModel:
         lazy_separate=lazy_separate,
         lazy_enforce=lazy_enforce,
     )
+
+
+def build_tsp_model_pyomo(
+    data: Union[str, TravelingSalesmanData],
+    solver: str = "gurobi_persistent",
+    params: Optional[dict[str, Any]] = None,
+) -> PyomoModel:
+
+    model = pe.ConcreteModel()
+    data = _tsp_read(data)
+
+    edges = tuplelist(
+        (i, j) for i in range(data.n_cities) for j in range(i + 1, data.n_cities)
+    )
+
+    # Decision variables
+    model.x = pe.Var(edges, domain=pe.Boolean, name="x")
+    model.obj = pe.Objective(
+        expr=sum(model.x[i, j] * data.distances[i, j] for (i, j) in edges)
+    )
+
+    # Eq: Must choose two edges adjacent to each node
+    model.degree_eqs = pe.ConstraintList()
+    for i in range(data.n_cities):
+        model.degree_eqs.add(
+            sum(model.x[min(i, j), max(i, j)] for j in range(data.n_cities) if i != j)
+            == 2
+        )
+
+    # Eq: Subtour elimination
+    model.subtour_eqs = pe.ConstraintList()
+
+    def lazy_separate(m: PyomoModel) -> List[Any]:
+        violations = []
+        m.solver.cbGetSolution([model.x[e] for e in edges])
+        x_val = {e: model.x[e].value for e in edges}
+        selected_edges = [e for e in edges if x_val[e] > 0.5]
+        graph = nx.Graph()
+        graph.add_edges_from(selected_edges)
+        for component in list(nx.connected_components(graph)):
+            if len(component) < data.n_cities:
+                cut_edges = tuple(
+                    (e[0], e[1])
+                    for e in edges
+                    if (e[0] in component and e[1] not in component)
+                    or (e[0] not in component and e[1] in component)
+                )
+                violations.append(cut_edges)
+        return violations
+
+    def lazy_enforce(m: PyomoModel, violations: List[Any]) -> None:
+        logger.warning(f"Adding {len(violations)} subtour elimination constraints...")
+        for violation in violations:
+            m.add_constr(
+                model.subtour_eqs.add(sum(model.x[e[0], e[1]] for e in violation) >= 2)
+            )
+
+    pm = PyomoModel(
+        model,
+        solver,
+        lazy_separate=lazy_separate,
+        lazy_enforce=lazy_enforce,
+    )
+    _pyomo_set_params(pm, params, solver)
+    return pm
+
+
+def _tsp_read(data: Union[str, TravelingSalesmanData]) -> TravelingSalesmanData:
+    if isinstance(data, str):
+        data = read_pkl_gz(data)
+    assert isinstance(data, TravelingSalesmanData)
+    return data

miplearn/solvers/pyomo.py

@@ -53,7 +53,12 @@ class PyomoModel(AbstractModel):
         assert (
             self.solver_name == "gurobi_persistent"
         ), "Callbacks are currently only supported on gurobi_persistent"
-        _gurobi_add_constr(self.solver, self.where, constr)
+        if self.where in [AbstractModel.WHERE_CUTS, AbstractModel.WHERE_LAZY]:
+            _gurobi_add_constr(self.solver, self.where, constr)
+        else:
+            # outside callbacks, add_constr shouldn't do anything, as the constraint
+            # has already been added to the ConstraintList object
+            pass
 
     def add_constrs(
         self,