This commit is contained in:
Marcin Kostrzewski 2022-06-16 21:05:08 +02:00
parent 14495cebda
commit e9d501c132

View File

@ -13,7 +13,7 @@ class Node:
return self.id, self.is_infected return self.id, self.is_infected
def __repr__(self): def __repr__(self):
return f"id: {self.id}, infected: {self.is_infected}" return f'id: {self.id}, infected: {self.is_infected}'
class Edge: class Edge:
@ -23,7 +23,7 @@ class Edge:
self.weight = weight self.weight = weight
def as_tuple(self): def as_tuple(self):
return self.node_a, self.node_b, {"weight": self.weight} return self.node_a, self.node_b, {'weight': self.weight}
class Graph: class Graph:
@ -53,8 +53,19 @@ def update(num, layout, g_repr, ax, our_graph: Graph):
for n in our_graph.get_nodes(): for n in our_graph.get_nodes():
n.is_infected = bool(random.getrandbits(1)) n.is_infected = bool(random.getrandbits(1))
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]
nx.draw_networkx(g_repr, ax=ax, pos=layout, node_color=colors, with_labels=False) sizes = [50 if n.is_infected else 1 for n in g_repr]
nx.draw(
g_repr,
ax=ax,
pos=layout,
node_color=colors,
with_labels=False,
node_size=sizes,
node_shape="s",
alpha=0.5,
linewidths=40,
)
def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int): def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int):
@ -62,13 +73,18 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int)
# Convert our graph class into tuples understood by networkx # Convert our graph class into tuples understood by networkx
g_repr.add_edges_from([e.as_tuple() for e in in_graph.edges]) g_repr.add_edges_from([e.as_tuple() for e in in_graph.edges])
layout = nx.spring_layout(g_repr) layout = nx.spring_layout(g_repr, k=0.3)
fig, ax = plt.subplots() fig, ax = plt.subplots()
fig.set_figwidth(15)
fig.set_figheight(15)
anim = animation.FuncAnimation( anim = animation.FuncAnimation(
fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph) fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph)
) )
anim.save(output_file_name) anim.save(output_file_name)
plt.style.use('seaborn')
plt.show() plt.show()
@ -107,13 +123,13 @@ def main():
] ]
) )
do_graph_animation("test.gif", network, 5) do_graph_animation('test.gif', network, 5)
bus = bus_network() bus = bus_network()
do_graph_animation("bus.gif", bus, 5) do_graph_animation('bus.gif', bus, 5)
ring = ring_network() ring = ring_network()
do_graph_animation("ring.gif", ring, 5) do_graph_animation('ring.gif', ring, 5)
if __name__ == "__main__": if __name__ == "__main__":