| Conditions | 12 |
| Total Lines | 157 |
| Code Lines | 81 |
| Lines | 0 |
| Ratio | 0 % |
| Changes | 0 | ||
Small methods make your code easier to understand, in particular if combined with a good name. Besides, if your method is small, finding a good name is usually much easier.
For example, if you find yourself adding comments to a method's body, this is usually a good sign to extract the commented part to a new method, and use the comment as a starting point when coming up with a good name for this new method.
Commonly applied refactorings include:
If many parameters/temporary variables are present:
Complex classes like amd._network_simplex.network_simplex() often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.
Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.
| 1 | """An implementation of the Wasserstein metric between two compositional vectors. |
||
| 12 | @numba.njit(cache=True) |
||
| 13 | def network_simplex(source_demands, sink_demands, network_costs): |
||
| 14 | """ |
||
| 15 | This is a port of the network simplex algorithm implented by Loïc Séguin-C |
||
| 16 | for the networkx package to allow acceleration via the numba package. |
||
| 17 | Copyright (C) 2010 Loïc Séguin-C. [email protected] |
||
| 18 | All rights reserved. |
||
| 19 | BSD license. |
||
| 20 | |||
| 21 | References |
||
| 22 | ---------- |
||
| 23 | [1] Z. Kiraly, P. Kovacs. |
||
| 24 | Efficient implementation of minimum-cost flow algorithms. |
||
| 25 | Acta Universitatis Sapientiae, Informatica 4(1):67--118. 2012. |
||
| 26 | [2] R. Barr, F. Glover, D. Klingman. |
||
| 27 | Enhancement of spanning tree labeling procedures for network |
||
| 28 | optimization. |
||
| 29 | INFOR 17(1):16--34. 1979. |
||
| 30 | """ |
||
| 31 | |||
| 32 | n_sources, n_sinks = source_demands.shape[0], sink_demands.shape[0] |
||
| 33 | network_costs = network_costs.ravel() |
||
| 34 | n = n_sources + n_sinks |
||
| 35 | e = n_sources * n_sinks |
||
| 36 | B = np.int64(np.ceil(np.sqrt(e))) |
||
| 37 | fp_multiplier = np.int64(1_000_000) |
||
| 38 | |||
| 39 | # Add one additional node for a dummy source and sink |
||
| 40 | source_d_int = (source_demands * fp_multiplier).astype(np.int64) |
||
| 41 | sink_d_int = (sink_demands * fp_multiplier).astype(np.int64) |
||
| 42 | |||
| 43 | # FP conversion error correction |
||
| 44 | source_sum = np.sum(source_d_int) |
||
| 45 | sink_sum = np.sum(sink_d_int) |
||
| 46 | sink_source_sum_diff = sink_sum - source_sum |
||
| 47 | |||
| 48 | if sink_source_sum_diff > 0: |
||
| 49 | source_d_int[np.argmax(source_d_int)] += sink_source_sum_diff |
||
| 50 | elif sink_source_sum_diff < 0: |
||
| 51 | sink_d_int[np.argmax(sink_d_int)] -= sink_source_sum_diff |
||
| 52 | |||
| 53 | demands = np.empty(n, dtype=np.int64) |
||
| 54 | demands[:n_sources] = -source_d_int |
||
| 55 | demands[n_sources:] = sink_d_int |
||
| 56 | |||
| 57 | tails = np.empty(e + n, dtype=np.int64) |
||
| 58 | heads = np.empty(e + n, dtype=np.int64) |
||
| 59 | |||
| 60 | for i in range(n_sources): |
||
| 61 | for j in range(n_sinks): |
||
| 62 | ind = i * n_sinks + j |
||
| 63 | tails[ind] = i |
||
| 64 | heads[ind] = j + n_sources |
||
| 65 | |||
|
|
|||
| 66 | for i, demand in enumerate(demands): |
||
| 67 | if demand > 0: |
||
| 68 | tails[e + i] = -1 |
||
| 69 | heads[e + i] = -1 |
||
| 70 | else: |
||
| 71 | tails[e + i] = i |
||
| 72 | heads[e + i] = i |
||
| 73 | |||
| 74 | # Create costs and capacities for the arcs between nodes |
||
| 75 | network_costs = network_costs * fp_multiplier |
||
| 76 | |||
| 77 | network_capac = np.array([min(source_demands[i], sink_demands[j]) |
||
| 78 | for i in range(n_sources) |
||
| 79 | for j in range(n_sinks)], |
||
| 80 | dtype=np.float64) * fp_multiplier |
||
| 81 | |||
| 82 | faux_inf = 3 * np.max(np.array(( |
||
| 83 | np.sum(network_capac), |
||
| 84 | np.sum(np.abs(network_costs)), |
||
| 85 | np.amax(source_d_int), |
||
| 86 | np.amax(sink_d_int)), |
||
| 87 | dtype=np.int64)) |
||
| 88 | |||
| 89 | # allocate arrays |
||
| 90 | costs = np.empty(e + n, dtype=np.int64) |
||
| 91 | costs[:e] = network_costs |
||
| 92 | costs[e:] = faux_inf |
||
| 93 | |||
| 94 | capac = np.empty(e + n, dtype=np.int64) |
||
| 95 | capac[:e] = network_capac |
||
| 96 | capac[e:] = fp_multiplier |
||
| 97 | |||
| 98 | flows = np.empty(e + n, dtype=np.int64) |
||
| 99 | flows[:e] = 0 |
||
| 100 | flows[e:e+n_sources] = source_d_int |
||
| 101 | flows[e+n_sources:] = sink_d_int |
||
| 102 | |||
| 103 | potentials = np.empty(n, dtype=np.int64) |
||
| 104 | demands_neg_mask = demands <= 0 |
||
| 105 | potentials[demands_neg_mask] = faux_inf |
||
| 106 | potentials[~demands_neg_mask] = -faux_inf |
||
| 107 | |||
| 108 | parent = np.empty(n + 1, dtype=np.int64) |
||
| 109 | parent[:-1] = -1 |
||
| 110 | parent[-1] = -2 |
||
| 111 | |||
| 112 | size = np.empty(n + 1, dtype=np.int64) |
||
| 113 | size[:-1] = 1 |
||
| 114 | size[-1] = n + 1 |
||
| 115 | |||
| 116 | next_node = np.arange(1, n + 2, dtype=np.int64) |
||
| 117 | next_node[-2] = -1 |
||
| 118 | next_node[-1] = 0 |
||
| 119 | |||
| 120 | last_node = np.arange(n + 1, dtype=np.int64) |
||
| 121 | last_node[-1] = n - 1 |
||
| 122 | |||
| 123 | prev_node = np.arange(-1, n, dtype=np.int64) |
||
| 124 | edge = np.arange(e, e + n, dtype=np.int64) |
||
| 125 | |||
| 126 | ### Pivot loop ### |
||
| 127 | |||
| 128 | f = 0 |
||
| 129 | while True: |
||
| 130 | i, p, q, f = find_entering_edges(B, e, f, tails, heads, costs, potentials, flows) |
||
| 131 | if p == -1: # If no entering edges then the optimal score is found |
||
| 132 | break |
||
| 133 | |||
| 134 | cycle_nodes, cycle_edges = find_cycle(i, p, q, size, edge, parent) |
||
| 135 | j, s, t = find_leaving_edge(cycle_nodes, cycle_edges, capac, flows, tails, heads) |
||
| 136 | augment_flow(cycle_nodes, cycle_edges, residual_capacity(j, s, capac, flows, tails), tails, flows) |
||
| 137 | |||
| 138 | if i != j: # Do nothing more if the entering edge is the same as the leaving edge. |
||
| 139 | if parent[t] != s: |
||
| 140 | # Ensure that s is the parent of t. |
||
| 141 | s, t = t, s |
||
| 142 | |||
| 143 | # Ensure that q is in the subtree rooted at t. |
||
| 144 | for val in cycle_edges: |
||
| 145 | if val == j: |
||
| 146 | p, q = q, p |
||
| 147 | break |
||
| 148 | elif val == i: |
||
| 149 | break |
||
| 150 | |||
| 151 | remove_edge(s, t, size, prev_node, last_node, next_node, parent, edge) |
||
| 152 | make_root(q, parent, size, last_node, prev_node, next_node, edge) |
||
| 153 | add_edge(i, p, q, next_node, prev_node, last_node, size, parent, edge) |
||
| 154 | update_potentials(i, p, q, heads, potentials, costs, last_node, next_node) |
||
| 155 | |||
| 156 | final_flows = flows[:e] / fp_multiplier |
||
| 157 | edge_costs = costs[:e] / fp_multiplier |
||
| 158 | final = final_flows.dot(edge_costs) |
||
| 159 | |||
| 160 | return final, final_flows.reshape((n_sources, n_sinks)) |
||
| 161 | |||
| 162 | |||
| 163 | @numba.njit(cache=True) |
||
| 164 | def reduced_cost(i, costs, potentials, tails, heads, flows): |
||
| 165 | """Return the reduced cost of an edge i. |
||
| 166 | """ |
||
| 167 | c = costs[i] - potentials[tails[i]] + potentials[heads[i]] |
||
| 168 | |||
| 169 | if flows[i] == 0: |
||
| 455 |