# Penning Trap Simulation: src/scripts/plot_2_particles.py Source File

Loading...
Searching...
No Matches
plot_2_particles.py
Go to the documentation of this file.
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
sns.set_theme()
params = {
"font.family": "Serif",
"font.serif": "Roman",
"text.usetex": True,
"axes.titlesize": "large",
"axes.labelsize": "large",
"xtick.labelsize": "large",
"ytick.labelsize": "large",
"legend.fontsize": "medium"
}
plt.rcParams.update(params)
def main():
directories = {
"output/simulate_2_particles/no_interaction/",
"output/simulate_2_particles/with_interaction/",
}
files = [
"particle_0_r.txt",
"particle_1_r.txt",
]
outputs = [
"../latex/images/simulate_2_particles_no_interaction_xy.pdf",
"../latex/images/simulate_2_particles_interaction_xy.pdf",
]
labels = [
r"$p_1$",
r"$p_2$",
r"$\hat{p}_1$",
r"$\hat{p}_2$",
]
colors = [
"seagreen",
"darkred"
]
start_pos = set()
for dir, output in zip(directories, outputs):
fig, ax = plt.subplots()
for label, color, file in zip(labels, colors, files):
with open(dir+file) as f:
lines = f.readlines()
t = np.linspace(0, 50, len(lines))
r = np.array([list(map(float, line.strip().split(","))) for line in lines])
ax.plot(r[:,0], r[:,1], label=label, color=color)
start_pos.add((r[0,0],r[0,1]))
for pos in start_pos:
ax.plot(*pos, "o", color="black")
ax.set_xlabel(r"x $(\mu m)$")
ax.set_ylabel(r"y $(\mu m)$")
ax.legend(loc="upper right")
ax.axis("equal")
fig.savefig(output, bbox_inches="tight")
if __name__ == "__main__":
main()