Added outline for nodes under attack in a given step.

This commit is contained in:
Krystian Wasilewski 2022-06-20 19:15:24 +02:00
parent 53e8b1f414
commit 6269168fb8

View File

@ -10,6 +10,7 @@ class Node:
def __init__(self, is_infected=False): def __init__(self, is_infected=False):
self.id = random.randint(1, 2000000) self.id = random.randint(1, 2000000)
self.is_infected = is_infected self.is_infected = is_infected
self.under_attack = False
def as_tuple(self): def as_tuple(self):
return self.id, self.is_infected return self.id, self.is_infected
@ -50,6 +51,11 @@ class Graph:
nodes.add(edge.node_b) nodes.add(edge.node_b)
return nodes 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)]: def get_adjacent_nodes(self, node: Node) -> [(Node, int)]:
""" """
:param node: Node to search for :param node: Node to search for
@ -84,6 +90,7 @@ class Graph:
else: else:
to_be_infected = random.choice([n[0] for n in adjacent_nodes]) to_be_infected = random.choice([n[0] for n in adjacent_nodes])
to_be_infected.is_infected = True to_be_infected.is_infected = True
to_be_infected.under_attack = True
self.update_survived() 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) ax.set_title(f'Step: {num}', loc='right', fontsize=30)
colors = ['red' if n.is_infected else 'blue' for n in g_repr] 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( nx.draw(
g_repr, g_repr,
ax=ax, ax=ax,
pos=layout, pos=layout,
node_color=colors, node_color=colors,
linewidths=linewidths,
edgecolors=edgecolors,
with_labels=False, with_labels=False,
node_size=sizes, node_size=sizes,
alpha=0.7, alpha=0.7,
) )
our_graph.clear_attacked()
def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout): def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout):
g_repr = nx.Graph() g_repr = nx.Graph()