diff --git a/network_attack_propagation.py b/network_attack_propagation.py index 658fee9..4410634 100644 --- a/network_attack_propagation.py +++ b/network_attack_propagation.py @@ -10,6 +10,7 @@ class Node: def __init__(self, is_infected=False): self.id = random.randint(1, 2000000) self.is_infected = is_infected + self.under_attack = False def as_tuple(self): return self.id, self.is_infected @@ -50,6 +51,11 @@ class Graph: nodes.add(edge.node_b) return nodes + def clear_attacked(self) -> None: + for edge in self.edges: + edge.node_a.under_attack = False + edge.node_b.under_attack = False + def get_adjacent_nodes(self, node: Node) -> [(Node, int)]: """ :param node: Node to search for @@ -84,6 +90,7 @@ class Graph: else: to_be_infected = random.choice([n[0] for n in adjacent_nodes]) to_be_infected.is_infected = True + to_be_infected.under_attack = True self.update_survived() @@ -101,17 +108,23 @@ def update(num, layout, g_repr, ax, our_graph: Graph): ax.set_title(f'Step: {num}', loc='right', fontsize=30) colors = ['red' if n.is_infected else 'blue' for n in g_repr] - sizes = [200 if n.is_infected else 50 for n in g_repr] + edgecolors = ['black' if n.under_attack else 'none' for n in g_repr] + linewidths = [3 if c == 'black' else 0 for c in edgecolors] + sizes = [300 if n.is_infected else 150 for n in g_repr] nx.draw( g_repr, ax=ax, pos=layout, node_color=colors, + linewidths=linewidths, + edgecolors=edgecolors, with_labels=False, node_size=sizes, alpha=0.7, ) + our_graph.clear_attacked() + def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout): g_repr = nx.Graph()