diff --git a/network_attack_propagation.py b/network_attack_propagation.py index 41a9fcd..507183f 100644 --- a/network_attack_propagation.py +++ b/network_attack_propagation.py @@ -13,7 +13,7 @@ class Node: return self.id, self.is_infected def __repr__(self): - return f'id: {self.id}, infected: {self.is_infected}' + return f"id: {self.id}, infected: {self.is_infected}" class Edge: @@ -23,7 +23,7 @@ class Edge: self.weight = weight 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: @@ -53,7 +53,7 @@ def update(num, layout, g_repr, ax, our_graph: Graph): 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] + 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) @@ -65,7 +65,9 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int) layout = nx.spring_layout(g_repr) fig, ax = plt.subplots() - anim = animation.FuncAnimation(fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph)) + anim = animation.FuncAnimation( + fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph) + ) anim.save(output_file_name) plt.show() @@ -73,7 +75,18 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int) def bus_network(n=30) -> Graph: network = Graph() nodes = [Node() for _ in range(n)] - edges = [Edge(nodes[i], nodes[i+1], 1.0) for i in range(n-1)] + edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)] + + network.add_edges(edges) + return network + + +def ring_network(n=30) -> Graph: + network = Graph() + nodes = [Node() for _ in range(n)] + edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)] + end_edge = Edge(nodes[n - 1], nodes[0], 1.0) + edges.append(end_edge) network.add_edges(edges) return network @@ -83,19 +96,24 @@ def main(): network = Graph() nodes = [Node(True), Node(), Node(), Node(True), Node()] - network.add_edges([ - Edge(nodes[1], nodes[0], 0.02), - Edge(nodes[1], nodes[2], 0.2), - Edge(nodes[2], nodes[0], 0.7), - Edge(nodes[3], nodes[2], 0.2), - Edge(nodes[3], nodes[1], 0.2), - Edge(nodes[4], nodes[3], 0.2) - ]) + network.add_edges( + [ + Edge(nodes[1], nodes[0], 0.02), + Edge(nodes[1], nodes[2], 0.2), + Edge(nodes[2], nodes[0], 0.7), + Edge(nodes[3], nodes[2], 0.2), + Edge(nodes[3], nodes[1], 0.2), + Edge(nodes[4], nodes[3], 0.2), + ] + ) - do_graph_animation('test.gif', network, 5) + do_graph_animation("test.gif", network, 5) bus = bus_network() - do_graph_animation('bus.gif', bus, 5) + do_graph_animation("bus.gif", bus, 5) + + ring = ring_network() + do_graph_animation("ring.gif", ring, 5) if __name__ == "__main__":