Penning Trap Simulation Simulate particle behavior inside a Penning Trap

Loading...
Searching...
No Matches
animate_100_particles.py
1 import matplotlib.pyplot as plt
2 import numpy as np
3 from matplotlib import animation
4 from mpl_toolkits.mplot3d import Axes3D
5
6
7 def get_data(files):
8 res = []
9 for file in files:
10 arr = [[], [], []]
11 with open(file, encoding="utf8") as f:
12 lines = f.readlines()
13
14 for line in lines:
15 xi, yi, zi = map(float, line.strip().split(","))
16 arr[0].append(xi)
17 arr[1].append(yi)
18 arr[2].append(zi)
19 res.append(arr)
20
21 return np.array(res)
22
23
24 def update(num, lines, arr):
25 for line, a in zip(lines, arr):
26 line.set_data(a[:2, num - 1 : num])
27 line.set_3d_properties(a[2, num])
28
29
30 def animate():
31 plt.style.use("dark_background")
32 fig = plt.figure()
33 ax = fig.add_subplot(projection="3d")
34
35 arr = get_data([f"output/simulate_100_particles/particle_{i}.txt" for i in range(100)])
36
37 arr = arr[:, :, ::10]
38
39 N = len(arr[0][0])
40
41 lines = [ax.plot(*a[:, 1], "o")[0] for a in arr]
42
43 ax.set_title("100 particles inside a Penning trap")
44 plt.figtext(
45 0.5,
46 0.01,
47 "100 randomly generated particles "
48 "evolving over a time of 50 microseconds.",
49 fontsize=12,
50 horizontalalignment="center",
51 )
52
53 ax.set_xlim3d([-500.0, 500.0])
54 ax.set_xlabel("X (micrometers)")
55
56 ax.set_ylim3d([-500.0, 500.0])
57 ax.set_ylabel("Y (micrometers)")
58
59 ax.set_zlim3d([-500.0, 500.0])
60 ax.set_zlabel("Z (micrometers)")
61
62 ani = animation.FuncAnimation(
63 fig, update, N, fargs=(lines, arr), interval=1, blit=False
64 )
65
66 # ani.save("../images/100_particles.gif", writer=animation.FFMpegFileWriter(fps=50))
67 plt.show()
68
69
70 if __name__ == "__main__":
71 animate()