-
Notifications
You must be signed in to change notification settings - Fork 1
/
minCostFlow.py
175 lines (148 loc) · 6.25 KB
/
minCostFlow.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
"""
Originally written by Anthony Gitter
Edited by Chris Magnano for batch use
This script attempts to connect sources to targets in a graph
using a minimum-cost flow algorithm. Minimum cost flow is solved using
Google's OR-Tools library: https://developers.google.com/optimization/flow/mincostflow
updated to use python 3
"""
import argparse
from ortools.graph.python.min_cost_flow import SimpleMinCostFlow
def parse_nodes(node_file):
''' Parse a list of sources or targets and return a set '''
with open(node_file) as node_f:
lines = node_f.readlines()
nodes = set(map(str.strip, lines))
return nodes
def construct_digraph(edges_file, cap):
''' Parse a list of weighted undirected edges. Construct a weighted
directed graph in which an undirected edge is represented with a pair of
directed edges. Use the specified weight as the edge weight and a default
capacity of 1.
'''
G = SimpleMinCostFlow()
idDict = dict() #Hold names to number ids
curID = 0
default_capacity = int(cap)
with open(edges_file) as edges_f:
for line in edges_f:
tokens = line.strip().split()
node1 = tokens[0]
if not node1 in idDict:
idDict[node1] = curID
curID += 1
node2 = tokens[1]
if not node2 in idDict:
idDict[node2] = curID
curID += 1
#Google's solver can only handle int weights, so round to the 100th
w = int((1-(float(tokens[2])))*100)
G.add_arc_with_capacity_and_unit_cost(idDict[node1],idDict[node2], default_capacity, int(w))
G.add_arc_with_capacity_and_unit_cost(idDict[node2],idDict[node1], default_capacity, int(w))
idDict["maxID"] = curID
return G,idDict
def print_graph(graph):
''' Print the edges in a graph '''
print('\n'.join(sorted(map(str, graph.edges(data=True)))))
def add_sources_targets(G, sources, targets, idDict, flow):
''' Similar to ResponseNet, add an artificial source node that is connected
to the real source nodes with directed edges. Unlike ResponseNet, these
directed edges should have weight of 0 and Infinite capacity. Also add an
artificial target node that has directed edges from the real target nodes
with the same weights and capacities as the source node edges. The new
nodes must be named "source" and "target".
'''
default_weight = 0
default_capacity = flow*10
curID = idDict["maxID"]
idDict["source"] = curID
curID += 1
idDict["target"] = curID
for source in sources:
if source in idDict:
G.add_arc_with_capacity_and_unit_cost(idDict["source"],idDict[source], default_capacity, default_weight)
for target in targets:
if target in idDict:
G.add_arc_with_capacity_and_unit_cost(idDict[target],idDict["target"], default_capacity, default_weight)
def write_output_to_sif(G,out_file_name,idDict):
''' Convert a flow dictionary from networkx.min_cost_flow into a list
of directed edges with the flow. Edges are represented as tuples.
'''
out_file = open(out_file_name,"w")
names = {v: k for k, v in idDict.items()}
numE = 0
for i in range(G.num_arcs()):
node1 = names[G.head(i)]
node2 = names[G.tail(i)]
flow = G.flow(i)
if flow <= 0:
continue
if node1 in ["source","target"]:
continue
if node2 in ["source","target"]:
continue
numE+=1
out_file.write(node1+"\t"+node2+"\n")
print("Final network had %d edges" % numE)
out_file.close()
return
def min_cost_flow(G, flow, output, idDict):
''' Use the min cost flow algorithm to distribute the specified amount
of flow from sources to targets. The artificial source should have
demand = -flow and the traget should have demand = flow. output is the
filename of the output file. The graph should have artificial nodes
named "source" and "target".
'''
G.set_node_supply(idDict['source'],int(flow))
G.set_node_supply(idDict['target'],int(-1*flow))
print("Computing min cost flow")
if G.solve() == G.OPTIMAL:
print("Solved!")
else:
print("There was an issue with the solver")
return
write_output_to_sif(G,output,idDict)
def main(args):
''' Parse a weighted edge list, source list, and target list. Run
min cost flow or k-shortest paths on the graph to find source-target
paths. Write the solutions to a file.
'''
flow = int(args.flow)
sources = parse_nodes(args.sources_file)
targets = parse_nodes(args.targets_file)
G,idDict = construct_digraph(args.edges_file, args.capacity)
add_sources_targets(G, sources, targets, idDict, flow)
out_file = args.output+"_flow"+str(flow)+"_c"+str(args.capacity)+".sif"
min_cost_flow(G, flow, out_file, idDict)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--edges_file',
help='Network file. File should be in SIF file format.',
type=str,
required=True)
parser.add_argument('--sources_file',
help='File which denotes source nodes, with one node per line.',
type=str,
required=True)
parser.add_argument('--targets_file',
help='File which denotes source nodes, with one node per line.',
type=str,
required=True)
parser.add_argument('--flow',
help='The amount of flow pushed through the network.',
type=int,
required=False,
default=1)
parser.add_argument('--output',
help='Prefix for all output files.',
type=str,
required=True)
parser.add_argument('--capacity',
help='The amount of flow which can pass through a single edge.',
type=float,
required=False,
default=1.0)
args = parser.parse_args()
main(args)