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