animate_100_particles.py

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
 
 
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 - 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/particle_{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_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 (micrometers)")
 
    ax.set_ylim3d([-500.0, 500.0])
    ax.set_ylabel("Y (micrometers)")
 
    ax.set_zlim3d([-500.0, 500.0])
    ax.set_zlabel("Z (micrometers)")
 
    ani = animation.FuncAnimation(
        fig, update, N, fargs=(lines, arr), interval=1, blit=False
    )
 
    # ani.save("../images/100_particles.gif", writer=animation.FFMpegFileWriter(fps=50))
    plt.show()
 
 
if __name__ == "__main__":
    animate()