diff options
| -rw-r--r-- | schroedinger/schrodinger_plot.py | 30 |
1 files 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) |