Enable multiple edges in the same direction between two nodes (#157)

Handle multiple edges in input data for pandas and networkx. --------- Co-authored-by: torressa <[email protected]> Co-authored-by: Simon Bowly <[email protected]>
Gurobi · Jul 17, 2024 · 5effe40 · 5effe40
1 parent 7b0bf32
commit 5effe40
Show file tree

Hide file tree

Showing 5 changed files with 246 additions and 157 deletions.
diff --git a/src/gurobi_optimods/datasets.py b/src/gurobi_optimods/datasets.py
@@ -137,14 +137,21 @@ def load_portfolio():
 
 
 def _convert_pandas_to_digraph(
-    edge_data, node_data, capacity=True, cost=True, demand=True
+    edge_data,
+    node_data,
+    capacity=True,
+    cost=True,
+    demand=True,
+    use_multigraph=False,
 ):
     """
-    Convert from a pandas DataFrame to a networkx.DiGraph with the appropriate
+    Convert from a pandas DataFrame to a networkx.MultiDiGraph with the appropriate
     attributes. For edges: `capacity`, and `cost`. For nodes: `demand`.
     """
+    graph_type = nx.MultiDiGraph if use_multigraph else nx.DiGraph
+
     G = nx.from_pandas_edgelist(
-        edge_data.reset_index(), create_using=nx.DiGraph(), edge_attr=True
+        edge_data.reset_index(), create_using=graph_type(), edge_attr=True
     )
     if demand:
         for i, d in node_data.iterrows():

diff --git a/src/gurobi_optimods/min_cost_flow.py b/src/gurobi_optimods/min_cost_flow.py
@@ -61,14 +61,17 @@ def min_cost_flow_pandas(
     with create_env() as env, gp.Model(env=env) as model:
         model.ModelSense = GRB.MINIMIZE
 
+        source_label, target_label = arc_data.index.names
+
+        arc_data = arc_data.reset_index()
+
         arc_df = arc_data.gppd.add_vars(model, ub="capacity", obj="cost", name="flow")
 
-        source_label, target_label = arc_data.index.names
         balance_df = (
             pd.DataFrame(
                 {
-                    "inflow": arc_df["flow"].groupby(target_label).sum(),
-                    "outflow": arc_df["flow"].groupby(source_label).sum(),
+                    "inflow": arc_df.groupby(target_label)["flow"].sum(),
+                    "outflow": arc_df.groupby(source_label)["flow"].sum(),
                     "demand": demand_data["demand"],
                 }
             )
@@ -84,6 +87,7 @@ def min_cost_flow_pandas(
         if model.Status in [GRB.INFEASIBLE, GRB.INF_OR_UNBD]:
             raise ValueError("Unsatisfiable flows")
 
+        arc_df = arc_df.set_index([source_label, target_label])
         return model.ObjVal, arc_df["flow"].gppd.X
 
 
@@ -171,27 +175,34 @@ def min_cost_flow_networkx(G, *, create_env):
     logger.info(
         f"Solving min-cost flow with {len(G.nodes)} nodes and {len(G.edges)} edges"
     )
-
     with create_env() as env, gp.Model(env=env) as model:
+        use_multigraph = isinstance(G, nx.MultiGraph)
+
+        G = nx.MultiDiGraph(G)
+
         edges, capacities, costs = gp.multidict(
-            {(i, j): [d["capacity"], d["cost"]] for i, j, d in G.edges(data=True)}
+            {
+                (i, j, idx): [d["capacity"], d["cost"]]
+                for i, j, idx, d in G.edges(data=True, keys=True)
+            }
         )
+
         nodes = list(G.nodes(data=True))
         x = {
-            (i, j): model.addVar(
-                name=f"flow[{i},{j}]",
-                ub=capacities[i, j],
-                obj=costs[i, j],
+            (i, j, idx): model.addVar(
+                name=f"flow[{i},{j},{idx}]",
+                ub=capacities[i, j, idx],
+                obj=costs[i, j, idx],
             )
-            for i, j in edges
+            for i, j, idx in edges
         }
 
         flow_constrs = {}
         for n, data in nodes:
             flow_constrs[n] = model.addConstr(
                 (
-                    gp.quicksum(x[j, n] for j in G.predecessors(n))
-                    - gp.quicksum(x[n, j] for j in G.successors(n))
+                    gp.quicksum(x[ie] for ie in G.in_edges(n, keys=True))
+                    - gp.quicksum(x[oe] for oe in G.out_edges(n, keys=True))
                     == data["demand"]
                 ),
                 name=f"flow_balance[{n}]",
@@ -203,7 +214,7 @@ def min_cost_flow_networkx(G, *, create_env):
             raise ValueError("Unsatisfiable flows")
 
         # Create a new Graph with selected edges in the matching
-        resulting_flow = nx.DiGraph()
+        resulting_flow = nx.MultiDiGraph() if use_multigraph else nx.DiGraph()
         resulting_flow.add_nodes_from(nodes)
         resulting_flow.add_edges_from(
             [(edge[0], edge[1], {"flow": v.X}) for edge, v in x.items() if v.X > 0.1]

diff --git a/tests/test_graph_utils.py b/tests/test_graph_utils.py
@@ -3,25 +3,24 @@
 
 def check_solution_pandas(solution, candidates):
     # Checks whether the solution (`pd.Series`) matches any of the list of
-    # candidates (containing `dict`)
-    if any(solution.to_dict() == c for c in candidates):
-        return True
-    return False
+    # candidates. Each candidate is a list of tuples ((i, j), v) tuples,
+    # compare with the solution in sorted order.
+    solution_list = sorted(solution.items())
+    return any(solution_list == sorted(candidate) for candidate in candidates)
 
 
 def check_solution_scipy(solution, candidates):
     # Checks whether the solution (`sp.sparray`) matches any of the list of
     # candidates (containing `np.ndarray`)
     arr = solution.toarray()
-    if any(np.array_equal(arr, c) for c in candidates):
-        return True
-    return False
+    return any(np.array_equal(arr, c) for c in candidates)
 
 
 def check_solution_networkx(solution, candidates):
     # Checks whether the solution (`nx.DiGraph`) matches any of the list of
     # candidates (containing tuples dict `{(i, j): data}`)
-    sol_dict = {(i, j): d for i, j, d in solution.edges(data=True)}
-    if any(sol_dict == c for c in candidates):
-        return True
-    return False
+    solution_list = sorted(
+        [((i, j), data["flow"]) for i, j, data in solution.edges(data=True)],
+        key=str,
+    )
+    return any(solution_list == sorted(candidate, key=str) for candidate in candidates)
diff --git a/tests/test_max_flow.py b/tests/test_max_flow.py
@@ -31,23 +31,23 @@ def test_pandas(self):
         obj, sol = max_flow(edge_data, 0, 5)
         sol = sol[sol > 0]
         self.assertEqual(obj, self.expected_max_flow)
-        candidate = {
-            (0, 1): 1.0,
-            (0, 2): 2.0,
-            (1, 3): 1.0,
-            (2, 3): 1.0,
-            (2, 4): 1.0,
-            (3, 5): 2.0,
-            (4, 5): 1.0,
-        }
-        candidate2 = {
-            (0, 1): 1.0,
-            (0, 2): 2.0,
-            (1, 3): 1.0,
-            (2, 4): 2.0,
-            (3, 5): 1.0,
-            (4, 5): 2.0,
-        }
+        candidate = [
+            ((0, 1), 1.0),
+            ((0, 2), 2.0),
+            ((1, 3), 1.0),
+            ((2, 3), 1.0),
+            ((2, 4), 1.0),
+            ((3, 5), 2.0),
+            ((4, 5), 1.0),
+        ]
+        candidate2 = [
+            ((0, 1), 1.0),
+            ((0, 2), 2.0),
+            ((1, 3), 1.0),
+            ((2, 4), 2.0),
+            ((3, 5), 1.0),
+            ((4, 5), 2.0),
+        ]
         self.assertTrue(check_solution_pandas(sol, [candidate, candidate2]))
 
     def test_empty_pandas(self):
@@ -84,23 +84,23 @@ def test_networkx(self):
         G = datasets.simple_graph_networkx()
         obj, sol = max_flow(G, 0, 5)
         self.assertEqual(obj, self.expected_max_flow)
-        candidate = {
-            (0, 1): {"flow": 1},
-            (0, 2): {"flow": 2},
-            (1, 3): {"flow": 1},
-            (2, 4): {"flow": 2},
-            (3, 5): {"flow": 1},
-            (4, 5): {"flow": 2},
-        }
-        candidate2 = {
-            (0, 1): {"flow": 1.0},
-            (0, 2): {"flow": 2.0},
-            (1, 3): {"flow": 1.0},
-            (2, 3): {"flow": 1.0},
-            (2, 4): {"flow": 1.0},
-            (3, 5): {"flow": 2.0},
-            (4, 5): {"flow": 1.0},
-        }
+        candidate = [
+            ((0, 1), 1),
+            ((0, 2), 2),
+            ((1, 3), 1),
+            ((2, 4), 2),
+            ((3, 5), 1),
+            ((4, 5), 2),
+        ]
+        candidate2 = [
+            ((0, 1), 1.0),
+            ((0, 2), 2.0),
+            ((1, 3), 1.0),
+            ((2, 3), 1.0),
+            ((2, 4), 1.0),
+            ((3, 5), 2.0),
+            ((4, 5), 1.0),
+        ]
         self.assertTrue(check_solution_networkx(sol, [candidate, candidate2]))
 
     @unittest.skipIf(nx is None, "networkx is not installed")
@@ -120,16 +120,16 @@ def test_pandas(self):
         obj, sol = max_flow(edge_data, 0, 4)
         sol = sol[sol > 0]
         self.assertEqual(obj, self.expected_max_flow)
-        candidate = {
-            (0, 1): 15.0,
-            (0, 2): 8.0,
-            (1, 3): 4.0,
-            (1, 2): 1.0,
-            (1, 4): 10.0,
-            (2, 3): 4.0,
-            (2, 4): 5.0,
-            (3, 4): 8.0,
-        }
+        candidate = [
+            ((0, 1), 15.0),
+            ((0, 2), 8.0),
+            ((1, 3), 4.0),
+            ((1, 2), 1.0),
+            ((1, 4), 10.0),
+            ((2, 3), 4.0),
+            ((2, 4), 5.0),
+            ((3, 4), 8.0),
+        ]
         self.assertTrue(check_solution_pandas(sol, [candidate]))
 
     def test_scipy(self):
@@ -152,23 +152,23 @@ def test_networkx(self):
         G = load_graph2_networkx()
         obj, sol = max_flow(G, 0, 4)
         self.assertEqual(obj, self.expected_max_flow)
-        candidate = {
-            (0, 1): {"flow": 15.0},
-            (0, 2): {"flow": 8.0},
-            (1, 3): {"flow": 4.0},
-            (1, 2): {"flow": 1.0},
-            (1, 4): {"flow": 10.0},
-            (2, 3): {"flow": 4.0},
-            (2, 4): {"flow": 5.0},
-            (3, 4): {"flow": 8.0},
-        }
-        candidate2 = {
-            (0, 1): {"flow": 15},
-            (0, 2): {"flow": 8},
-            (1, 2): {"flow": 1},
-            (1, 3): {"flow": 4},
-            (1, 4): {"flow": 10},
-            (2, 3): {"flow": 9},
-            (3, 4): {"flow": 13},
-        }
+        candidate = [
+            ((0, 1), 15.0),
+            ((0, 2), 8.0),
+            ((1, 3), 4.0),
+            ((1, 2), 1.0),
+            ((1, 4), 10.0),
+            ((2, 3), 4.0),
+            ((2, 4), 5.0),
+            ((3, 4), 8.0),
+        ]
+        candidate2 = [
+            ((0, 1), 15),
+            ((0, 2), 8),
+            ((1, 2), 1),
+            ((1, 3), 4),
+            ((1, 4), 10),
+            ((2, 3), 9),
+            ((3, 4), 13),
+        ]
         self.assertTrue(check_solution_networkx(sol, [candidate, candidate2]))