Merge branch 'plot'

1 files changed, 27 insertions, 15 deletions

diff --git a/schroedinger/schrodinger_plot.py b/schroedinger/schrodinger_plot.py
index ba16655..f9e48b6 100644
--- a/schroedinger/schrodinger_plot.py
+++ b/schroedinger/schrodinger_plot.py
@@ -1,8 +1,9 @@
+"""Plots the solution of schroedinger/schrodinger_solve.py"""
 import sys
 import argparse
+from pathlib import Path
 import numpy as np
 import matplotlib.pyplot as plt
-from pathlib import Path
 
 sys.path.insert(0, str(Path.cwd().parent))
 
@@ -25,6 +26,7 @@ def plot_wavefuncs(ax: plt.Axes, solution_path: Path, wavefunction_scale: float
     :param ax: The axis on which the potential should be drawn
     :param solution_path: The complete path where the solution is stored
     :param wavefunction_scale: Scales the wavefuncs-data before the energy offset is applied
+    :return: (min_value, max_value), used for scaling of the y-axis
     """
     wavefuncs_data = np.loadtxt(f'{solution_path}/wavefuncs.dat')
     energies = np.loadtxt(f'{solution_path}/energies.dat')
@@ -39,7 +41,7 @@ def plot_wavefuncs(ax: plt.Axes, solution_path: Path, wavefunction_scale: float
     energy_delta = np.abs(energies.max() - energies.min()) / energies.shape[0] * 2
     max_value = energies.max() + energy_delta
     min_value = energies.min() - energy_delta
-    return max_value, min_value
+    return min_value, max_value
 
 
 
@@ -87,16 +89,19 @@ def plot_solution(
         uncertainty_lim: None | tuple[float, float]
 ):
     """
-    Main function for plotting the solution of schrodinger_solve.py. By default, the plot is shown directly and also
-    saved to the location of the solution.
-    :param solution_dir: if None the current directory is used. If str the complete path to the solution directory.
-    :param output_dir: if None the current directory is used. If str the complete path to the output directory.
+    Main function for plotting the solution of schrodinger_solve.py.
+     By default, the plot is shown directly and also saved to the location of the solution.
+    :param solution_dir: if None the current directory is used.
+     If str the complete path to the solution directory.
+    :param output_dir: if None the current directory is used.
+     If str the complete path to the output directory.
     :param show_plot: if True the plot is shown in the default viewer.
     :param export_pdf: if True the plot is exported as a pdf in the selected directory.
     :param wavefunction_scale: scales the wavefunc-data before the energy is added.
     :param energy_lim: limits the y-axis of both axis. Needs to be given as (min, max).
     :param x_lim: limits the x-axis of the left axis. Needs to be given as (min, max).
-    :param uncertainty_lim: limits the x-axis of the right axis. Needs to be given as (min, max)
+    :param uncertainty_lim: limits the x-axis of the right axis.
+     Needs to be given as (min, max)
     """
     solution_path = Path(solution_dir) if solution_dir else Path.cwd()
 
@@ -105,15 +110,16 @@ def plot_solution(
     ax2: plt.Axes = fig.add_subplot(122)
 
     plot_potential(ax1, solution_path)
-    max_value, min_value = plot_wavefuncs(ax1, solution_path, wavefunction_scale)
+    automatic_limit = plot_wavefuncs(ax1, solution_path, wavefunction_scale)
     plot_expected_value(ax1, ax2, solution_path)
 
     if energy_lim is not None:
         y_lim = energy_lim
     else:
-        y_lim = min_value, max_value
+        y_lim = automatic_limit
 
-    ax1.set(xlabel='x [Bohr]', ylabel='Energy [Hartree]', title=r'Potential, Eigenstate, $\langle x \rangle$',
+    ax1.set(xlabel='x [Bohr]', ylabel='Energy [Hartree]',
+            title=r'Potential, Eigenstate, $\langle x \rangle$',
             xlim=x_lim, ylim=y_lim)
     ax2.set(xlabel='[Bohr]', title=r'$\sigma _{x}$', yticks=[], xlim=uncertainty_lim, ylim=y_lim)
     if export_pdf:
@@ -125,6 +131,9 @@ def plot_solution(
 
 
 def main():
+    """
+    Main call of function.
+    """
     # argparse setup
     parser = argparse.ArgumentParser(description='Plots the solutions from schrodinger_solve.py')
     msg = 'the path of the solution directory (default: None)'
@@ -137,17 +146,20 @@ def main():
     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)'
+    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)'
+    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)'
+    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()
 
     # main call to plot the solution
-    plot_solution(args.solution_dir, args.output_dir, args.show, args.export, args.scale, args.energy_lim, args.xlim,
-                  args.uncertainty_lim)
+    plot_solution(args.solution_dir, args.output_dir, args.show, args.export, args.scale,
+                  args.energy_lim, args.xlim, args.uncertainty_lim)
 
 
 if __name__ == '__main__':