Colorful nodes

network_attack_propagation.py

@@ -36,13 +36,25 @@ class Graph:
     def add_edges(self, edges: [Edge]):
         [self.edges.append(e) for e in edges]
 
+    def get_nodes(self) -> [Node]:
+        nodes = set()
+        for edge in self.edges:
+            nodes.add(edge.node_a)
+            nodes.add(edge.node_b)
+        return nodes
 
-def update(num, layout, g_repr, ax, our_graph):
+
+def update(num, layout, g_repr, ax, our_graph: 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)
+
+    for n in our_graph.get_nodes():
+        n.is_infected = bool(random.getrandbits(1))
+
+    colors = ['red' if n.is_infected else 'blue' for n in g_repr]
+    nx.draw_networkx(g_repr, ax=ax, pos=layout, node_color=colors, with_labels=False)
 
 
 def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int):
@@ -71,7 +83,7 @@ def main():
         Edge(nodes[4], nodes[3], 0.2)
     ])
 
-    do_graph_animation('test.gif', network, 1)
+    do_graph_animation('test.gif', network, 5)
 
 
 if __name__ == "__main__":