Graph-Algorithms/helper.py at main · osherboudara99/Graph-Algorithms · GitHub

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
import os
import networkx as nx
import matplotlib.pyplot as plt
import imageio
from IPython.display import Image


def clear_marks(G):
    for node in G.nodes:
        G.nodes[node]['marked'] = False
    for edge in G.edges:
        G.edges[edge]['marked'] = False

def draw_graph(G, with_labels=True, with_tree=True, ax=None):

    marked_nodes = [node for node, attrs in G.nodes.items() if attrs.get('marked', False) == True]
    unmarked_nodes = list(set(G.nodes) - set(marked_nodes))

    marked_edges = [e for e, attrs in G.edges.items() if attrs.get('marked', False) == True]
    unmarked_edges = list(set(G.edges) - set(marked_edges))

    pos = nx.spring_layout(G) # positions for all nodes
    if with_tree:
        nx.draw_networkx_nodes(G, pos,
                           nodelist=marked_nodes,
                           node_color='#FF00FF',
                           node_size=400,
                           ax=ax)

        nx.draw_networkx_edges(G, pos,
                            edgelist=marked_edges,
                            edge_color='#FF00FF',
                            width=2,
                            arrows=True,
                            ax=ax)

        nx.draw_networkx_nodes(G, pos,
                           nodelist=unmarked_nodes,
                           node_color='#7EFF67',
                           node_size=400,
                           ax=ax)

        nx.draw_networkx_edges(G, pos,
                               edgelist=unmarked_edges,
                               edge_color='k',
                               width=1,
                               arrows=True,
                               ax=ax)
    else:
        nx.draw_networkx_nodes(G, pos,
                           nodelist=unmarked_nodes + marked_nodes,
                           node_color='#7EFF67',
                           node_size=400,
                           ax=ax)

        nx.draw_networkx_edges(G, pos,
                               edgelist=unmarked_edges + marked_edges,
                               edge_color='k',
                               width=1,
                               arrows=True,
                               ax=ax)

    if with_labels:
        labels = {n: n for n in G.nodes}
        nx.draw_networkx_labels(G, pos, labels, font_size=14, ax=ax)

    if ax:
        ax.axis('off')
    plt.axis('off')


def animate_graph(G, edges, name):
    pos = nx.spring_layout(G)
    fig, ax = plt.subplots()

    nx.draw_networkx_nodes(G, pos,
                           nodelist=G.nodes,
                           node_color='#00bbff',
                           ax=ax)

    nx.draw_networkx_edges(G, pos,
                           edgelist=G.edges,
                           arrows=True,
                           ax=ax)
    labels = {n: n for n in G.nodes}
    nx.draw_networkx_labels(G, pos, labels, font_size=14, ax=ax)
    ax.axis('off')
    plt.savefig('.tmp0.png')
    tmp_files = ['.tmp0.png']
    marked_edges = []
    for i, edge in enumerate(edges, 1):
        marked_edges.append(edge)
        if isinstance(G, nx.DiGraph):
            unmarked_edges = set(G.edges) - set(marked_edges)
        else:
            unmarked_edges = set(G.edges) - set(marked_edges) - set(e[::-1] for e in marked_edges)
        ax.clear()
        nx.draw_networkx_nodes(G, pos,
                               nodelist=G.nodes,
                               node_color='#00bbff',
                               ax=ax)
        nx.draw_networkx_edges(G, pos,
                               edgelist=marked_edges,
                               edge_color ='#b300f0',
                               arrows=True,
                               ax=ax)
        nx.draw_networkx_edges(G, pos,
                               edgelist=unmarked_edges,
                               arrows=True,
                               ax=ax)
        nx.draw_networkx_labels(G, pos, labels, font_size=14, ax=ax)
        ax.axis('off')
        tmp_files.append(f'.tmp{i}.png')
        plt.savefig(tmp_files[-1])
    ax.remove()
    with imageio.get_writer(f'static/{name}', mode='I') as writer:
        for file in tmp_files:
            image = imageio.imread(file)
            for _ in range(5):
                writer.append_data(image)
            os.remove(file)
        for _ in range(10):
            writer.append_data(image)

    return Image(url=f'static/{name}')