Penning Trap Simulation Simulate particle behavior inside a Penning Trap

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