aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--schroedinger/schrodinger_plot.py30
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)