Make lazy constr component compatible with Pyomo+Gurobi

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