From 12da0fe191d88cdf76788a834d6e3d2edae7655f Mon Sep 17 00:00:00 2001 From: Marcin Kostrzewski Date: Thu, 16 Jun 2022 20:14:56 +0200 Subject: [PATCH] Colorful nodes --- network_attack_propagation.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/network_attack_propagation.py b/network_attack_propagation.py index fa7fc5a..1135094 100644 --- a/network_attack_propagation.py +++ b/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__":