Project-3/src/scripts/plot_3d.py
2023-10-24 21:49:26 +02:00

68 lines
1.8 KiB
Python

## @file plot_3d.py
#
# @author Cory Alexander Balaton (coryab)
# @author Janita Ovidie Sandtrøen Willumsen (janitaws)
#
# @version 1.0
#
# @brief Plot 2 particles with and without particle interactions in 3D.
#
# @bug No known bugs
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",
"figure.autolayout": True
}
plt.rcParams.update(params)
def main():
files = [
"output/simulate_2_particles/no_interaction/particle_0_r.txt",
"output/simulate_2_particles/no_interaction/particle_1_r.txt",
"output/simulate_2_particles/with_interaction/particle_0_r.txt",
"output/simulate_2_particles/with_interaction/particle_1_r.txt"
]
labels = [
r"$p_1$",
r"$p_2$",
r"$\hat{p}_1$",
r"$\hat{p}_2$",
]
colors = [
"mediumaquamarine",
"salmon",
"seagreen",
"darkred"
]
ax = plt.figure().add_subplot(projection="3d")
for label, color, file in zip(labels, colors, files):
with open(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], r[:,2], label=label, color=color)
ax.set_xlabel(r"x $(\mu m)$")
ax.set_ylabel(r"y $(\mu m)$")
ax.set_zlabel(r"z $(\mu m)$")
ax.view_init(10,-35)
plt.legend()
plt.savefig("../latex/images/3d_plot.pdf")
if __name__ == "__main__":
main()