From 8373deea81125611251d855b732359d28c358b39 Mon Sep 17 00:00:00 2001
From: Cory <cory@balaton.dev>
Date: Sun, 22 Oct 2023 16:34:48 +0200
Subject: [PATCH] Update plots

---
 src/scripts/plot_phase_space.py     | 63 ++++++++++++++++-------------
 src/scripts/plot_single_particle.py |  6 +--
 2 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/src/scripts/plot_phase_space.py b/src/scripts/plot_phase_space.py
index 849dbfe..b989961 100644
--- a/src/scripts/plot_phase_space.py
+++ b/src/scripts/plot_phase_space.py
@@ -20,21 +20,13 @@ def main():
         "output/simulate_2_particles/no_interaction/",
         "output/simulate_2_particles/with_interaction/",
     }
-    titles = {
-        "particles without interaction",
-        "particles with interaction"
-    }
     files = [
-        "particle_0_r.txt",
-        "particle_0_v.txt",
-        "particle_1_r.txt",
-        "particle_1_v.txt",
+        "particle_0",
+        "particle_1",
     ]
     labels = [
-        r"particle 1 r",
-        r"particle 1 v",
-        r"particle 2 r",
-        r"particle 2 v",
+        [r"$p_1$", r"$\hat{p}_1$"],
+        [r"$p_2$", r"$\hat{p}_2$"]
     ] 
     colors = [
         "lightskyblue",
@@ -42,25 +34,42 @@ def main():
         "salmon",
         "tomato",
     ]
-    fig1, axs1 = plt.subplots(2,1)
-    fig2, axs2 = plt.subplots(2,1)
-    for i, (dir, title) in enumerate(zip(directories, titles)):
-        for label, color, file in zip(labels, colors, files):
-            with open(dir+file) as f:
+    linestyles = [
+        "solid",
+        "dotted"
+    ]
+    fig1, axs1 = plt.subplots(2,1,sharex=True)
+    fig2, axs2 = plt.subplots(2,1,sharex=True)
+    for i, (file, label) in enumerate(zip(files, labels)):
+        for j, (dir, linestyle) in enumerate(zip(directories, linestyles)):
+            with open(dir+file+"_r.txt") 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])
-                axs1[i].plot(t, r[:,0], label=label, color=color)
-                axs2[i].plot(t, r[:,2], label=label, color=color)
-        
-            axs1[i].set(xlabel=r"t $(\mu s)$", ylabel = r"z $(\mu m)$")
-        
-        axs1[i].legend()
-        axs1[i].set_title(title)
+                axs1[i].plot(t, r[:,0], label=f"{label[j]} r", linestyle=linestyle)
+                axs2[i].plot(t, r[:,2], label=f"{label[j]} r", linestyle=linestyle)
 
-    fig1.savefig("../latex/images/phase_space_2_particles_x.pdf")
-    fig2.savefig("../latex/images/phase_space_2_particles_z.pdf")
-    plt.show()
+            with open(dir+file+"_v.txt") 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])
+                axs1[i].plot(t, r[:,0], label=f"{label[j]} v", linestyle=linestyle)
+                axs2[i].plot(t, r[:,2], label=f"{label[j]} v", linestyle=linestyle)
+        
+        
+        axs1[1].set(xlabel=r"t $(\mu s)$", ylabel = r"x $(\mu m)$")
+        axs1[0].set(ylabel = r"x $(\mu m)$")
+        axs2[1].set(xlabel=r"t $(\mu s)$", ylabel = r"z $(\mu m)$")
+        axs2[0].set(ylabel = r"z $(\mu m)$")
+
+        axs1[i].legend()
+        # axs1[i].set_title(title)
+        axs2[i].legend()
+        # axs2[i].set_title(title)
+
+    fig1.savefig("../latex/images/phase_space_2_particles_x.pdf", bbox_inches="tight")
+    fig2.savefig("../latex/images/phase_space_2_particles_z.pdf", bbox_inches="tight")
+    # plt.show()
 
 
 if __name__ == "__main__":
diff --git a/src/scripts/plot_single_particle.py b/src/scripts/plot_single_particle.py
index 6ee2d16..a80a8b4 100644
--- a/src/scripts/plot_single_particle.py
+++ b/src/scripts/plot_single_particle.py
@@ -39,9 +39,9 @@ def main():
         r = np.array([list(map(float, line.strip().split(","))) for line in lines])
 
     fig, ax = plt.subplots()
-    ax.plot(t, r[:, 2], label="approximation", color=colors[3])
-    ax.plot(t, z(t), label="analytical", color=colors[5], linestyle="dotted")
-    ax.set_xlabel(r"Time $(\mu s)$")
+    ax.plot(t, r[:, 2], label="approximation", color="mediumseagreen")
+    ax.plot(t, z(t), label="analytical", color="black", linestyle="dotted")
+    ax.set_xlabel(r"t $(\mu s)$")
     ax.set_xlim((-5, 55))
     ax.set_ylabel(r"z $(\mu m)$")
     ax.set_ylim((-25, 25))