command line functionality for plot

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)