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

Loading...
Searching...
No Matches
plot_phase_space.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",
"particle_1",
]
labels = [r"$p_1$", r"$p_2$"]
output = [
"../latex/images/phase_space_no_interaction",
"../latex/images/phase_space_interaction",
]
colors = [
"seagreen",
"darkred",
]
for i, dir in enumerate(directories):
fig1, axs1 = plt.subplots(sharex=True)
fig2, axs2 = plt.subplots(sharex=True)
for j, (file, label, color) in enumerate(zip(files, labels, colors)):
r = []
v = []
with open(dir + file + "_r.txt") as f:
lines = f.readlines()
r = np.array(
[list(map(float, line.strip().split(","))) for line in lines]
)
with open(dir + file + "_v.txt") as f:
lines = f.readlines()
v = np.array(
[list(map(float, line.strip().split(","))) for line in lines]
)
axs1.plot(r[:, 0], v[:, 0], label=label, color=color)
axs1.plot(r[0,0], v[0,0], "ko")
axs2.plot(r[:, 2], v[:, 2], label=label, color=color)
axs2.plot(r[0,2], v[0,2], "ko")
axs1.axis("equal")
axs2.axis("equal")
axs1.set(xlabel=r"$x (\mu m)$", ylabel=r"$v_x (\mu m / \mu s)$")
axs2.set(xlabel=r"$z (\mu m)$", ylabel=r"$v_z (\mu m / \mu s)$")
axs1.legend(loc="upper right")
axs2.legend(loc="upper right")
fig1.savefig(f"{output[i]}_x.pdf", bbox_inches="tight")
fig2.savefig(f"{output[i]}_z.pdf", bbox_inches="tight")
# plt.show()
if __name__ == "__main__":
main()