Project-3/src/animate_100_particles.py
2023-10-08 00:34:03 +02:00

65 lines
1.4 KiB
Python

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, encoding="utf8") as f:
lines = f.readlines()
for line in lines:
xi,yi,zi = map(float, 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():
plt.style.use("dark_background")
fig = plt.figure()
ax = fig.add_subplot(projection="3d")
arr = get_data([f"output/p{i}_RK4.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()