diff --git a/network_attack_propagation.py b/network_attack_propagation.py
index b385f57..0e1cad7 100644
--- a/network_attack_propagation.py
+++ b/network_attack_propagation.py
@@ -68,12 +68,12 @@ def update(num, layout, g_repr, ax, our_graph: Graph):
     )
 
 
-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, layout):
     g_repr = nx.Graph()
     # 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, k=0.3)
+    layout = layout(g_repr)
     fig, ax = plt.subplots()
 
     fig.set_figwidth(15)
@@ -123,13 +123,13 @@ def main():
         ]
     )
 
-    do_graph_animation('test.gif', network, 5)
+    do_graph_animation('test.gif', network, 5, nx.spring_layout)
 
     bus = bus_network()
-    do_graph_animation('bus.gif', bus, 5)
+    do_graph_animation('bus.gif', bus, 5, nx.spiral_layout)
 
     ring = ring_network()
-    do_graph_animation('ring.gif', ring, 5)
+    do_graph_animation('ring.gif', ring, 5, nx.circular_layout)
 
 
 if __name__ == "__main__":