Added dynamic graph updating & drawing

network_attack_propagation.py

@@ -1,6 +1,19 @@
+import random
+import networkx as nx
+import matplotlib.pyplot as plt
+from matplotlib import animation
+
+
 class Node:
-    def __init__(self):
-        self.is_infected = False
+    def __init__(self, is_infected=False):
+        self.id = random.randint(1, 2000000)
+        self.is_infected = is_infected
+
+    def as_tuple(self):
+        return self.id, self.is_infected
+
+    def __repr__(self):
+        return f'id: {self.id}, infected: {self.is_infected}'
 
 
 class Edge:
@@ -9,10 +22,57 @@ class Edge:
         self.node_b = node_b
         self.weight = weight
 
+    def as_tuple(self):
+        return self.node_a, self.node_b, {'weight': self.weight}
+
 
 class Graph:
     def __init__(self):
         self.edges = []
 
     def add_edge(self, edge: Edge):
-        self.edges.append(edge)
+        self.edges.append(edge)
+
+    def add_edges(self, edges: [Edge]):
+        [self.edges.append(e) for e in edges]
+
+
+def update(num, layout, g_repr, ax, our_graph):
+    """
+    This function is called every 'step', so if you wish to update the graph, do it here
+    """
+    ax.clear()
+    nx.draw_networkx(g_repr, ax=ax, pos=layout)
+
+
+def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int):
+    g_repr = nx.Graph()
+    # Convert our graph class into tuples understood by networkx
+    g_repr.add_edges_from([e.as_tuple() for e in in_graph.edges])
+
+    layout = nx.spring_layout(g_repr)
+    fig, ax = plt.subplots()
+
+    anim = animation.FuncAnimation(fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph))
+    anim.save(output_file_name)
+    plt.show()
+
+
+def main():
+    network = Graph()
+    nodes = [Node(True), Node(), Node(), Node(True), Node()]
+
+    network.add_edges([
+        Edge(nodes[1], nodes[0], 0.02),
+        Edge(nodes[1], nodes[2], 0.2),
+        Edge(nodes[2], nodes[0], 0.7),
+        Edge(nodes[3], nodes[2], 0.2),
+        Edge(nodes[3], nodes[1], 0.2),
+        Edge(nodes[4], nodes[3], 0.2)
+    ])
+
+    do_graph_animation('test.gif', network, 1)
+
+
+if __name__ == "__main__":
+    main()