From cad3708a0708410efe797e9e970329d62266b85c Mon Sep 17 00:00:00 2001
From: Cory <cory@balaton.dev>
Date: Tue, 26 Dec 2023 15:31:57 +0100
Subject: [PATCH] Add changes

---
 python_scripts/colormap.py     | 59 ++++++++++++++++++++++++++++------
 python_scripts/colormap_alt.py | 38 ----------------------
 2 files changed, 49 insertions(+), 48 deletions(-)
 delete mode 100644 python_scripts/colormap_alt.py

diff --git a/python_scripts/colormap.py b/python_scripts/colormap.py
index 1d4dd55..f596e15 100644
--- a/python_scripts/colormap.py
+++ b/python_scripts/colormap.py
@@ -1,6 +1,7 @@
-import numpy as np
-import matplotlib.pyplot as plt
 import ast
+
+import matplotlib.pyplot as plt
+import numpy as np
 import seaborn as sns
 
 sns.set_theme()
@@ -16,22 +17,60 @@ params = {
 }
 plt.rcParams.update(params)
 
+
 def plot():
     with open("data/color_map.txt") as f:
-        lines = f.readlines();
+        lines = f.readlines()
         size = int(lines[0])
         for i, line in enumerate(lines[1:]):
-            fig, ax = plt.subplots()
+            # Create figures for each plot
+            fig1, ax1 = plt.subplots()
+            fig2, ax2 = plt.subplots()
+            fig3, ax3 = plt.subplots()
+
             arr = line.strip().split("\t")
-            arr = np.asarray(list(map(lambda x: ((a := complex(*ast.literal_eval(x)))*a.conjugate()).real, arr)))
+            arr = np.asarray(list(map(lambda x: complex(*ast.literal_eval(x)), arr)))
+            # Reshape and transpose array
+            arr = arr.reshape(size, size).T
 
-            arr = arr.reshape(size,size)
+            # Plot color maps
+            color_map1 = ax1.imshow(
+                np.multiply(arr, arr.conj()).real,
+                interpolation="nearest",
+                cmap=sns.color_palette("mako", as_cmap=True),
+            )
+            color_map2 = ax2.imshow(
+                arr.real,
+                interpolation="nearest",
+                cmap=sns.color_palette("mako", as_cmap=True),
+            )
+            color_map3 = ax3.imshow(
+                arr.imag,
+                interpolation="nearest",
+                cmap=sns.color_palette("mako", as_cmap=True),
+            )
 
-            ax.imshow(arr.T, interpolation="nearest")
-            plt.savefig(f"latex/images/color_map_{i}.pdf")
-            plt.close(fig)
+            # Create color bar
+            fig1.colorbar(color_map1, ax=ax1)
+            fig2.colorbar(color_map2, ax=ax2)
+            fig3.colorbar(color_map3, ax=ax3)
+
+            # Remove grids
+            ax1.grid(False)
+            ax2.grid(False)
+            ax3.grid(False)
+
+            # Save the figures
+            fig1.savefig(f"latex/images/color_map_{i}_prob.pdf")
+            fig2.savefig(f"latex/images/color_map_{i}_real.pdf")
+            fig3.savefig(f"latex/images/color_map_{i}_imag.pdf")
+
+            # Close figures
+            plt.close(fig1)
+            plt.close(fig2)
+            plt.close(fig3)
 
 
 if __name__ == "__main__":
     plot()
-    
+
diff --git a/python_scripts/colormap_alt.py b/python_scripts/colormap_alt.py
deleted file mode 100644
index 1496a37..0000000
--- a/python_scripts/colormap_alt.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import numpy as np
-import matplotlib.pyplot as plt
-import ast
-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 plot():
-    with open("data/color_map.txt") as f:
-        lines = f.readlines()
-        size = int(lines[0])
-        for i, line in enumerate(lines[1:]):
-            fig, ax = plt.subplots()
-            arr = line.strip().split("\t")
-            arr = np.asarray(list(map(lambda x: ((a := complex(*ast.literal_eval(x)))*a.conjugate()).real, arr)))
-
-            arr = arr.reshape(size,size)
-            
-            color_map = ax.imshow(arr.T, interpolation="nearest", cmap=sns.color_palette("mako", as_cmap=True))
-            fig.colorbar(color_map, ax=ax)
-            plt.grid(False)
-            plt.savefig(f"latex/images/color_map_{i}.pdf")
-            plt.close(fig)
-
-
-if __name__ == "__main__":
-    plot()
\ No newline at end of file