Add descriptions

src/animate_100_particles.py

@@ -1,7 +1,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
-from mpl_toolkits.mplot3d import Axes3D
 from matplotlib import animation
+from mpl_toolkits.mplot3d import Axes3D
+
 
 def get_data(files):
     res = []
@@ -19,20 +20,19 @@ def get_data(files):
 
     return np.array(res)
 
+
 def update(num, lines, arr):
     for line, a in zip(lines, arr):
-        line.set_data(a[:2, num])
+        line.set_data(a[:2, num - 1 : num])
         line.set_3d_properties(a[2, num])
 
 
-
 def animate():
     plt.style.use("dark_background")
     fig = plt.figure()
     ax = fig.add_subplot(projection="3d")
 
-
+    arr = get_data([f"output/simulate_100_particles/p{i}.txt" for i in range(100)])
-    arr = get_data([f"output/p{i}_RK4.txt" for i in range(100)])
 
     arr = arr[:, :, ::10]
 
@@ -40,25 +40,32 @@ def animate():
 
     lines = [ax.plot(*a[:, 1], "o")[0] for a in arr]
 
+    ax.set_title("100 particles inside a Penning trap")
+    plt.figtext(
+        0.5,
+        0.01,
+        "100 randomly generated particles "
+        "evolving over a time of 50 microseconds.",
+        fontsize=12,
+        horizontalalignment="center",
+    )
+
     ax.set_xlim3d([-500.0, 500.0])
-    ax.set_xlabel('X')
+    ax.set_xlabel("X (micrometers)")
 
     ax.set_ylim3d([-500.0, 500.0])
-    ax.set_ylabel('Y')
+    ax.set_ylabel("Y (micrometers)")
 
     ax.set_zlim3d([-500.0, 500.0])
-    ax.set_zlabel('Z')  
+    ax.set_zlabel("Z (micrometers)")
 
-    ani = animation.FuncAnimation(fig, update, N, fargs=(lines, arr), 
+    ani = animation.FuncAnimation(
-                                  interval=1,
+        fig, update, N, fargs=(lines, arr), interval=1, blit=False
-                                  blit=False)
+    )
 
-
+    ani.save("../images/100_particles.gif", writer=animation.FFMpegFileWriter(fps=50))
-    ani.save("100_particles.gif", writer=animation.FFMpegFileWriter(fps=30))
+    # plt.show()
-    plt.show()
 
 
 if __name__ == "__main__":
     animate()
-
-