diff options
| -rw-r--r-- | schroedinger/schrodinger_plot.py | 42 |
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__': |