Styling

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,8 +53,19 @@ 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]
-    nx.draw_networkx(g_repr, ax=ax, pos=layout, node_color=colors, with_labels=False)
+    colors = ['red' if n.is_infected else 'blue' for n in g_repr]
+    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):
@@ -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
     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.set_figwidth(15)
+    fig.set_figheight(15)
+
     anim = animation.FuncAnimation(
         fig, update, frames=frame_count, fargs=(layout, g_repr, ax, in_graph)
     )
     anim.save(output_file_name)
+
+    plt.style.use('seaborn')
     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()
-    do_graph_animation("bus.gif", bus, 5)
+    do_graph_animation('bus.gif', bus, 5)
 
     ring = ring_network()
-    do_graph_animation("ring.gif", ring, 5)
+    do_graph_animation('ring.gif', ring, 5)
 
 
 if __name__ == "__main__":