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 = []
@@ -11,7 +12,7 @@ def get_data(files):
             lines = f.readlines()
 
             for line in lines:
-                xi,yi,zi = map(float, line.strip().split(","))
+                xi, yi, zi = map(float, line.strip().split(","))
                 arr[0].append(xi)
                 arr[1].append(yi)
                 arr[2].append(zi)
@@ -19,46 +20,52 @@ 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]
+    arr = arr[:, :, ::10]
 
     N = len(arr[0][0])
 
-    lines = [ax.plot(*a[:,1], "o")[0] for a in arr]
+    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), 
-                                  interval=1,
-                                  blit=False)
+    ani = animation.FuncAnimation(
+        fig, update, N, fargs=(lines, arr), interval=1, blit=False
+    )
 
-
-    ani.save("100_particles.gif", writer=animation.FFMpegFileWriter(fps=30))
-    plt.show()
+    ani.save("../images/100_particles.gif", writer=animation.FFMpegFileWriter(fps=50))
+    # plt.show()
 
 
 if __name__ == "__main__":
     animate()
-
-