From 12a631e40f3596c5b7b928cb3031563c9d4be7ca Mon Sep 17 00:00:00 2001 From: Tim Suhling Date: Fri, 12 Jul 2024 15:31:29 +0200 Subject: Plot functionality for test data --- schroedinger/schrodinger_plot.py | 74 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 73 insertions(+), 1 deletion(-) 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() -- cgit v1.2.3 From 093e00e42c00f92f88d05d9e05476d06807c485d Mon Sep 17 00:00:00 2001 From: Tim Suhling Date: Fri, 12 Jul 2024 16:06:04 +0200 Subject: command line functionality for plot --- schroedinger/schrodinger_plot.py | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/schroedinger/schrodinger_plot.py b/schroedinger/schrodinger_plot.py index f572ae6..48032d4 100644 --- a/schroedinger/schrodinger_plot.py +++ b/schroedinger/schrodinger_plot.py @@ -1,10 +1,32 @@ import sys +import argparse import numpy as np import matplotlib.pyplot as plt from pathlib import Path sys.path.insert(0, str(Path.cwd().parent)) +# argparse setup +parser = argparse.ArgumentParser(description='Plots the solutions from schrodinger_solve.py') +msg = 'the path of the solution directory (default: .)' +parser.add_argument('-s', '--solution', default='.', help=msg) +msg = 'the path where the pdf should be saved (default: .)' +parser.add_argument('-p', '--pdf', default='.', help=msg) +msg = 'Boolean, if True the plot is shown directly (default: True)' +parser.add_argument('--show', default=True, help=msg, type=bool) +msg = 'Boolean, if True the plot is exported as a pdf (default: True)' +parser.add_argument('-e', '--export', default=True, help=msg, type=bool) +msg = 'Float, scales the wave functions (default: 1.0)' +parser.add_argument('--scale', default=1.0, help=msg, type=float) +msg = 'Limit of the x-axis of the left plot. None or tuple[float, float] of shape (x_min, x_max)(default: None)' +parser.add_argument('-x', '--xlim', default=None, help=msg) +msg = 'Limit of the y-axis of the left plot. None or tuple[float, float] of shape (y_min, y_max)(default: None)' +parser.add_argument('-y1', '--energy_lim', default=None, help=msg) +msg = 'Limit of the y-axis of the right plot. None or tuple[float, float] of shape (y_min, y_max)(default: None)' +parser.add_argument('-y2', '--uncertainty_lim', default=None, help=msg) +args = parser.parse_args() + + def plot_potential(ax: plt.Axes, solution_path: str): potential_data = np.loadtxt(f'{solution_path}/potential.dat') x = potential_data[:, 0] @@ -48,9 +70,11 @@ def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str): 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): +def plot_solution(solution_path: str = args.solution, pdf_path: str = args.pdf, show_plot: bool = args.show, + export_pdf: bool = args.export, wavefunction_scale: float = args.scale, + energy_lim: None | tuple[float, float] = args.energy_lim, + x_lim: None | tuple[float, float] = args.xlim, + uncertainty_lim: None | tuple[float, float] = args.uncertainty_lim): 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) -- cgit v1.2.3 From 1c0aee1f87f9f35da422f51601e3baffda071d19 Mon Sep 17 00:00:00 2001 From: Tim Suhling Date: Fri, 19 Jul 2024 14:46:13 +0200 Subject: Small improvement to plot --- schroedinger/schrodinger_plot.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/schroedinger/schrodinger_plot.py b/schroedinger/schrodinger_plot.py index 48032d4..b9ad14f 100644 --- a/schroedinger/schrodinger_plot.py +++ b/schroedinger/schrodinger_plot.py @@ -30,8 +30,8 @@ args = parser.parse_args() 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='') + 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): @@ -51,8 +51,10 @@ def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str): expected_values = expvalues_data[:, 0] uncertainties = expvalues_data[:, 1] + uncertainty_max = uncertainties.max() - x_lim = ax1.get_xlim() + x1_lim = ax1.get_xlim() + x2_lim = (0, uncertainty_max*1.1) y_lim = ax1.get_ylim() for index in range(expvalues_data.shape[0]): energy = energies[index] @@ -60,14 +62,13 @@ def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str): 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) + ax1.hlines(energy, *x1_lim, colors='gray', alpha=0.5) - ax2.hlines(energy, *x_lim, colors='gray', alpha=0.5) + ax2.hlines(energy, *x2_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) + ax1.set(xlim=x1_lim, ylim=y_lim) + ax2.set(xlim=x2_lim, ylim=y_lim) def plot_solution(solution_path: str = args.solution, pdf_path: str = args.pdf, show_plot: bool = args.show, @@ -94,7 +95,7 @@ def plot_solution(solution_path: str = args.solution, pdf_path: str = args.pdf, def main(): - plot_solution('test', 'test') + plot_solution('test', 'test', x_lim=(-5, 5), energy_lim=(0, 3.5), wavefunction_scale=1.0) if __name__ == '__main__': -- cgit v1.2.3