Kolorowanie ścieżki podczas przepływu za pomocą modułu networkx

Question

Witam,

Czy jest jakiś sposób aby pokolorować ścieżkę na inny kolor niż z góry nałożony szary gdy występuje własnie na niej przepływ? Nie potrafię wyciągnąć obiektu(danej ścieżki) gdzie właśnie występuje przepływ, myślę że wtedy mógłbym zastosować jakiś warunek logiczny który by zmieniał kolor danej ścieżki przez co przepływ byłby bardziej czytelny. Dodatkowo wtedy spróbowałbym zmienić kolor ścieżki  na czerwony gdy osiągnęła by swoją maksymalną przepustowość.

 

import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch


def ford_fulkerson(graph, source, sink, debug=None):
    flow, path = 0, True

    while path:

        path, reserve = depth_first_search(graph, source, sink)
        flow += reserve


        for v, u in zip(path, path[1:]):
            if graph.has_edge(v, u):
                graph[v][u]['flow'] += reserve
            else:
                graph[u][v]['flow'] -= reserve


        if callable(debug):
            debug(graph, path, reserve, flow)


def depth_first_search(graph, source, sink):
    undirected = graph.to_undirected()
    explored = {source}
    stack = [(source, 0, dict(undirected[source]))]

    while stack:
        v, _, neighbours = stack[-1]
        if v == sink:
            break


        while neighbours:
            u, e = neighbours.popitem()
            if u not in explored:
                break
        else:
            stack.pop()
            continue


        in_direction = graph.has_edge(v, u)
        capacity = e['capacity']
        flow = e['flow']
        neighbours = dict(undirected[u])


        if in_direction and flow < capacity:
            stack.append((u, capacity - flow, neighbours))
            explored.add(u)
        elif not in_direction and flow:
            stack.append((u, flow, neighbours))
            explored.add(u)


    reserve = min((f for _, f, _ in stack[1:]), default=0)
    path = [v for v, _, _ in stack]

    return path, reserve


graph = nx.DiGraph()
graph.add_nodes_from('ABCDEFGH')
graph.add_edges_from([
    ('A', 'B', {'capacity': 4, 'flow': 0}),
    ('A', 'C', {'capacity': 5, 'flow': 0}),
    ('A', 'D', {'capacity': 7, 'flow': 0}),
    ('B', 'E', {'capacity': 7, 'flow': 0}),
    ('C', 'E', {'capacity': 6, 'flow': 0}),
    ('C', 'F', {'capacity': 4, 'flow': 0}),
    ('C', 'G', {'capacity': 1, 'flow': 0}),
    ('D', 'F', {'capacity': 8, 'flow': 0}),
    ('D', 'G', {'capacity': 1, 'flow': 0}),
    ('E', 'H', {'capacity': 7, 'flow': 0}),
    ('F', 'H', {'capacity': 6, 'flow': 0}),
    ('G', 'H', {'capacity': 4, 'flow': 0}),
])



layout = {
    'A': [0, 1], 'B': [1, 2], 'C': [1, 1], 'D': [1, 0],
    'E': [2, 2], 'F': [2, 1], 'G': [2, 0], 'H': [3, 1],
}


def draw_graph():
    plt.figure(figsize=(12, 4))
    plt.axis('off')
    nx.draw_networkx_nodes(graph, layout, node_color='steelblue', node_size=600)
    nx.draw_networkx_labels(graph, layout, font_color='white')
    nx.draw_networkx_edges(graph, layout, edge_color='gray')

    for u, v, e in graph.edges(data=True):
        edge_list = list(graph.edges)
        #print(edge_list)
        #print(len(graph))

        label = '{}/{}'.format(e['flow'], e['capacity'])
        color = 'green' if e['flow'] < e['capacity'] else 'red'
        #edge_color = 'yellow' if e['flow'] < e['capacity'] else 'red'
        #graph.update(edge_list[0], color='edge_color')
        #nx.set_edge_attributes(graph, "edge_color")
        #nx.draw_networkx_edges(graph, layout, color="edge_color")
        #nx.draw_networkx_edges(graph1, layout, edge_color='green')
        x = layout[u][0] * .6 + layout[v][0] * .4
        y = layout[u][1] * .6 + layout[v][1] * .4


        t = plt.text(x, y, label, size=16, color=color,
                     horizontalalignment='center', verticalalignment='center')

    plt.show()


draw_graph()


def flow_debug(graph, path, reserve, flow):
    print('flow increased by', reserve,
          'at path', path,
          '; current flow', flow)
    draw_graph()



ford_fulkerson(graph, 'A', 'H', flow_debug)