src/test.py

@@ -0,0 +1,63 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib import animation
+
+def get_data(files):
+    res = []
+    for file in files:
+        arr = [[], [], []]
+        with open(file) as f:
+            lines = f.readlines()
+
+            for line in lines:
+                xi,yi,zi = map(lambda x: float(x), line.strip().split(","))
+                arr[0].append(xi)
+                arr[1].append(yi)
+                arr[2].append(zi)
+        res.append(arr)
+
+    return np.array(res)
+
+def update(num, lines, arr):
+    for line, a in zip(lines, arr):
+        line.set_data(a[:2, num])
+        line.set_3d_properties(a[2, num])
+    
+
+
+def animate():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection="3d")
+
+
+    arr = get_data([f"output/p{i}.txt" for i in range(100)])
+
+    arr = arr[:,:,::10]
+
+    N = len(arr[0][0])
+
+    lines = [ax.plot(*a[:,1], "o")[0] for a in arr]
+
+    ax.set_xlim3d([-500.0, 500.0])
+    ax.set_xlabel('X')
+
+    ax.set_ylim3d([-500.0, 500.0])
+    ax.set_ylabel('Y')
+
+    ax.set_zlim3d([-500.0, 500.0])
+    ax.set_zlabel('Z')  
+
+    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()
+
+
+if __name__ == "__main__":
+    animate()
+
+