MIPLearn v0.3

miplearn/problems/tsp.py

@@ -1,22 +1,23 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 from dataclasses import dataclass
-from typing import List, Tuple, Any, Optional, Dict
+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 overrides import overrides
+from gurobipy import quicksum, GRB, tuplelist
 from scipy.spatial.distance import pdist, squareform
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
+import logging
 
-from miplearn.instance.base import Instance
-from miplearn.solvers.learning import InternalSolver
-from miplearn.solvers.pyomo.base import BasePyomoSolver
-from miplearn.types import ConstraintName
+from miplearn.io import read_pkl_gz
+from miplearn.solvers.gurobi import GurobiModel
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -25,80 +26,6 @@ class TravelingSalesmanData:
     distances: np.ndarray
 
 
-class TravelingSalesmanInstance(Instance):
-    """An instance ot the Traveling Salesman Problem.
-
-    Given a list of cities and the distance between each pair of cities, the problem
-    asks for the shortest route starting at the first city, visiting each other city
-    exactly once, then returning to the first city. This problem is a generalization
-    of the Hamiltonian path problem, one of Karp's 21 NP-complete problems.
-    """
-
-    def __init__(self, n_cities: int, distances: np.ndarray) -> None:
-        super().__init__()
-        assert isinstance(distances, np.ndarray)
-        assert distances.shape == (n_cities, n_cities)
-        self.n_cities = n_cities
-        self.distances = distances
-        self.edges = [
-            (i, j) for i in range(self.n_cities) for j in range(i + 1, self.n_cities)
-        ]
-
-    @overrides
-    def to_model(self) -> pe.ConcreteModel:
-        model = pe.ConcreteModel()
-        model.x = pe.Var(self.edges, domain=pe.Binary)
-        model.obj = pe.Objective(
-            expr=sum(model.x[i, j] * self.distances[i, j] for (i, j) in self.edges),
-            sense=pe.minimize,
-        )
-        model.eq_degree = pe.ConstraintList()
-        model.eq_subtour = pe.ConstraintList()
-        for i in range(self.n_cities):
-            model.eq_degree.add(
-                sum(
-                    model.x[min(i, j), max(i, j)]
-                    for j in range(self.n_cities)
-                    if i != j
-                )
-                == 2
-            )
-        return model
-
-    @overrides
-    def find_violated_lazy_constraints(
-        self,
-        solver: InternalSolver,
-        model: Any,
-    ) -> Dict[ConstraintName, List]:
-        selected_edges = [e for e in self.edges if model.x[e].value > 0.5]
-        graph = nx.Graph()
-        graph.add_edges_from(selected_edges)
-        violations = {}
-        for c in list(nx.connected_components(graph)):
-            if len(c) < self.n_cities:
-                cname = ("st[" + ",".join(map(str, c)) + "]").encode()
-                violations[cname] = list(c)
-        return violations
-
-    @overrides
-    def enforce_lazy_constraint(
-        self,
-        solver: InternalSolver,
-        model: Any,
-        component: List,
-    ) -> None:
-        assert isinstance(solver, BasePyomoSolver)
-        cut_edges = [
-            e
-            for e in self.edges
-            if (e[0] in component and e[1] not in component)
-            or (e[0] not in component and e[1] in component)
-        ]
-        constr = model.eq_subtour.add(expr=sum(model.x[e] for e in cut_edges) >= 2)
-        solver.add_constraint(constr)
-
-
 class TravelingSalesmanGenerator:
     """Random generator for the Traveling Salesman Problem."""
 
@@ -118,7 +45,7 @@ class TravelingSalesmanGenerator:
         distributions `n`, `x` and `y`. For each (unordered) pair of cities (i,j),
         the distance d[i,j] between them is set to:
 
-            d[i,j] = gamma[i,j] \sqrt{(x_i - x_j)^2 + (y_i - y_j)^2}
+            d[i,j] = gamma[i,j] \\sqrt{(x_i - x_j)^2 + (y_i - y_j)^2}
 
         where gamma is sampled from the provided probability distribution `gamma`.
 
@@ -183,3 +110,68 @@ class TravelingSalesmanGenerator:
         n = self.n.rvs()
         cities = np.array([(self.x.rvs(), self.y.rvs()) for _ in range(n)])
         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)
+
+    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")
+
+    model._x = x
+    model._edges = edges
+    model._n_cities = data.n_cities
+
+    # Objective function
+    model.setObjective(quicksum(x[(i, j)] * data.distances[i, j] for (i, j) in edges))
+
+    # Eq: Must choose two edges adjacent to each node
+    model.addConstrs(
+        (
+            quicksum(x[min(i, j), max(i, j)] for j in range(data.n_cities) if i != j)
+            == 2
+            for i in range(data.n_cities)
+        ),
+        name="eq_degree",
+    )
+
+    def find_violations(model: GurobiModel) -> List[Any]:
+        violations = []
+        x = model.inner.cbGetSolution(model.inner._x)
+        selected_edges = [e for e in model.inner._edges if x[e] > 0.5]
+        graph = nx.Graph()
+        graph.add_edges_from(selected_edges)
+        for component in list(nx.connected_components(graph)):
+            if len(component) < model.inner._n_cities:
+                cut_edges = [
+                    e
+                    for e in model.inner._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 fix_violations(model: GurobiModel, violations: List[Any], where: str) -> None:
+        for violation in violations:
+            constr = quicksum(model.inner._x[e[0], e[1]] for e in violation) >= 2
+            if where == "cb":
+                model.inner.cbLazy(constr)
+            else:
+                model.inner.addConstr(constr)
+        logger.info(f"tsp: added {len(violations)} subtour elimination constraints")
+
+    model.update()
+
+    return GurobiModel(
+        model,
+        find_violations=find_violations,
+        fix_violations=fix_violations,
+    )