From 14495cebda232a7026c6c8d1f739e813114939f6 Mon Sep 17 00:00:00 2001
From: Wirusik <wirus006@gmail.com>
Date: Thu, 16 Jun 2022 20:58:13 +0200
Subject: [PATCH] ring init

---
 network_attack_propagation.py | 48 ++++++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 15 deletions(-)

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__":