Plot functionality for test data

1 files changed, 73 insertions, 1 deletions

diff --git a/schroedinger/schrodinger_plot.py b/schroedinger/schrodinger_plot.py
index 9369e2e..f572ae6 100644
--- a/schroedinger/schrodinger_plot.py
+++ b/schroedinger/schrodinger_plot.py
@@ -1,5 +1,77 @@
+import sys
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+sys.path.insert(0, str(Path.cwd().parent))
+
+
+def plot_potential(ax: plt.Axes, solution_path: str):
+    potential_data = np.loadtxt(f'{solution_path}/potential.dat')
+    x = potential_data[:, 0]
+    V = potential_data[:, 1]
+    ax.plot(x, V, c='black', ls='-', marker='')
+
+
+def plot_wavefuncs(ax: plt.Axes, solution_path: str, wavefunction_scale: float = 1.0):
+    wavefuncs_data = np.loadtxt(f'{solution_path}/wavefuncs.dat')
+    energies = np.loadtxt(f'{solution_path}/energies.dat')
+    x = wavefuncs_data[:, 0]
+    colors = ['blue', 'red']
+    for wave_index in range(1, wavefuncs_data.shape[1]):
+        energy = energies[wave_index-1]
+        wave_function = wavefuncs_data[:, wave_index] * wavefunction_scale + energy
+        ax.plot(x, wave_function, color=colors[wave_index%2], zorder=100)
+
+
+def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str):
+    expvalues_data = np.loadtxt(f'{solution_path}/expvalues.dat')
+    energies = np.loadtxt(f'{solution_path}/energies.dat')
+
+    expected_values = expvalues_data[:, 0]
+    uncertainties = expvalues_data[:, 1]
+
+    x_lim = ax1.get_xlim()
+    y_lim = ax1.get_ylim()
+    for index in range(expvalues_data.shape[0]):
+        energy = energies[index]
+        expected_value = expected_values[index]
+        uncertainty = uncertainties[index]
+
+        ax1.plot(expected_value, energy, marker='x', ls='', c='green', zorder=101)
+        ax1.hlines(energy, *x_lim, colors='gray', alpha=0.5)
+
+        ax2.hlines(energy, *x_lim, colors='gray', alpha=0.5)
+        ax2.plot(uncertainty, energy, color='orange', marker='+', ls='', markersize=10)
+
+    uncertainty_max = uncertainties.max()
+    ax1.set(xlim=x_lim, ylim=y_lim)
+    ax2.set(xlim=(0, uncertainty_max*1.1), ylim=y_lim)
+
+
+def plot_solution(solution_path: str, pdf_path: str, show_plot: bool = True, export_pdf: bool = True,
+                  wavefunction_scale: float = 1.0, energy_lim: None | tuple[float, float] = None,
+                  x_lim: None | tuple[float, float] = None, uncertainty_lim: None | tuple[float, float] = None):
+    fig = plt.figure(dpi=200, figsize=(6, 4), tight_layout=True)
+    ax1: plt.Axes = fig.add_subplot(121)
+    ax2: plt.Axes = fig.add_subplot(122)
+
+    plot_potential(ax1, solution_path)
+    plot_wavefuncs(ax1, solution_path, wavefunction_scale)
+    plot_expected_value(ax1, ax2, solution_path)
+
+    ax1.set(xlabel='x [Bohr]', ylabel='Energy [Hartree]', title=r'Potential, Eigenstate, $\langle x \rangle$',
+            xlim=x_lim, ylim=energy_lim)
+    ax2.set(xlabel='[Bohr]', title=r'$\sigma _{x}$', yticks=[], xlim=uncertainty_lim, ylim=energy_lim)
+    if export_pdf:
+        plt.savefig(f'{pdf_path}/schroedinger_solution.pdf', dpi=300)
+    if show_plot:
+        plt.show()
+    plt.close()
+
+
 def main():
-    pass
+    plot_solution('test', 'test')
+
 
 if __name__ == '__main__':
     main()