From 8fe3c2bf8eaddc215d5aa634cc69a0d3d73cb9c2 Mon Sep 17 00:00:00 2001
From: Marcin Kostrzewski <komar86311@gmail.com>
Date: Tue, 21 Jun 2022 16:24:48 +0200
Subject: [PATCH] Refactor experiments

---
 network_attack_propagation.py | 54 +++++++++++++----------------------
 1 file changed, 20 insertions(+), 34 deletions(-)

diff --git a/network_attack_propagation.py b/network_attack_propagation.py
index f2bfa77..86ea46b 100644
--- a/network_attack_propagation.py
+++ b/network_attack_propagation.py
@@ -154,7 +154,7 @@ def degree_avg(edges, digits=2):
     return round(mean(degrees.values()), digits)
 
 
-def bus_network(n=30, infected_idx=0) -> tuple[Graph, float]:
+def bus_network(n=30, infected_idx=0) -> tuple[Graph, float, int]:
     network = Graph()
     nodes = [Node() for _ in range(n)]
     nodes[infected_idx].is_infected = True
@@ -181,7 +181,7 @@ def star_network(cluster_count=5, starsize=6) -> tuple[Graph, float, int]:
     return network, degree_avg(edges), node_count
 
 
-def ring_network(n=30) -> tuple[Graph, float]:
+def ring_network(n=30) -> tuple[Graph, float, int]:
     network = Graph()
     nodes = [Node() for _ in range(n)]
     nodes[0].is_infected = True
@@ -200,51 +200,37 @@ def summary(average_degrees: list[float], propagation_speeds: list[float]) -> No
     fig.savefig("summary.png")
     plt.show()
 
-def bus_experiment():
+
+def experiment(network, avg_degree, node_count):
     degrees = []
     speeds = []
     sizes = [0, 8, 15]
     for i in sizes:
-        bus, bus_avg_degree, node_count = bus_network(20 + i)
-        do_graph_animation(f'bus{i}.gif', bus, 90, nx.spiral_layout)
-        speeds.append(bus.rounds_survived / node_count)
-        degrees.append(bus_avg_degree)
+        do_graph_animation(f'bus{i}.gif', network, 90, nx.spiral_layout)
+        speeds.append(network.rounds_survived / node_count)
+        degrees.append(avg_degree)
         print(f"\n{node_count} NODE BUS")
-        print(f"average degree = {bus_avg_degree}")
+        print(f"average degree = {avg_degree}")
         print(f"propagation speed = {round(speeds[-1], 2)}")
-        print(f"bus_{i} rounds survived = {bus.rounds_survived + 1}")
+        print(f"bus_{i} rounds survived = {network.rounds_survived + 1}")
 
     summary(degrees, speeds)
 
+
+def bus_experiment():
+    bus, bus_avg_degree, node_count = bus_network(20 + i)
+    experiment(bus, bus_avg_degree, node_count)
+
+
 def ring_experiment():
-    degrees = []
-    speeds = []
-    sizes = [0, 8, 15]
-    for i in sizes:
-        ring, ring_avg_degree, node_count = ring_network(20 + i)
-        do_graph_animation(f'ring{i}.gif', ring, 90, nx.spiral_layout)
-        speeds.append(ring.rounds_survived / node_count)
-        degrees.append(ring_avg_degree)
-        print(f"\n{node_count} NODE RING")
-        print(f"average degree = {ring_avg_degree}")
-        print(f"propagation speed = {round(speeds[-1], 2)}")
-        print(f"ring_{i} rounds survived = {ring.rounds_survived + 1}")
+    ring, ring_avg_degree, node_count = ring_network(20 + i)
+    experiment(ring, ring_avg_degree, node_count)
+
 
 def star_experiment():
-    degrees = []
-    speeds = []
-    sizes = range(0, 8, 2)
-    for i in sizes:
-        star, star_avg_degree, node_count = star_network(cluster_count=3 + i, starsize=3 + i)
-        do_graph_animation(f'star{i}.gif', star, 120, nx.kamada_kawai_layout)
-        speeds.append(star.rounds_survived / node_count)
-        degrees.append(star_avg_degree)
-        print(f"\n{node_count} NODE STAR")
-        print(f"average degree = {star_avg_degree}")
-        print(f"propagation speed = {round(speeds[-1], 2)}")
-        print(f"star_{i} rounds survived = {star.rounds_survived + 1}")
+    star, star_avg_degree, node_count = star_network(cluster_count=3 + i, starsize=3 + i)
+    experiment(star, star_avg_degree, node_count)
 
-    summary(degrees, speeds)
 
 def main():
     bus_experiment()