1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
|
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]
v = potential_data[:, 1]
ax.plot(x, v, c='black', ls='-', marker='')
def plot_wavefuncs(ax: plt.Axes, solution_path: str, wavefunction_scale: float = 1.0):
wavefuncs_data = np.loadtxt(f'{solution_path}/wavefuncs.dat')
energies = np.loadtxt(f'{solution_path}/energies.dat')
x = wavefuncs_data[:, 0]
colors = ['blue', 'red']
for wave_index in range(1, wavefuncs_data.shape[1]):
energy = energies[wave_index-1]
wave_function = wavefuncs_data[:, wave_index] * wavefunction_scale + energy
ax.plot(x, wave_function, color=colors[wave_index%2], zorder=100)
def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: str):
expvalues_data = np.loadtxt(f'{solution_path}/expvalues.dat')
energies = np.loadtxt(f'{solution_path}/energies.dat')
expected_values = expvalues_data[:, 0]
uncertainties = expvalues_data[:, 1]
uncertainty_max = uncertainties.max()
x1_lim = ax1.get_xlim()
x2_lim = (0, uncertainty_max*1.1)
y_lim = ax1.get_ylim()
for index in range(expvalues_data.shape[0]):
energy = energies[index]
expected_value = expected_values[index]
uncertainty = uncertainties[index]
ax1.plot(expected_value, energy, marker='x', ls='', c='green', zorder=101)
ax1.hlines(energy, *x1_lim, colors='gray', alpha=0.5)
ax2.hlines(energy, *x2_lim, colors='gray', alpha=0.5)
ax2.plot(uncertainty, energy, color='orange', marker='+', ls='', markersize=10)
ax1.set(xlim=x1_lim, ylim=y_lim)
ax2.set(xlim=x2_lim, ylim=y_lim)
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)
plot_potential(ax1, solution_path)
plot_wavefuncs(ax1, solution_path, wavefunction_scale)
plot_expected_value(ax1, ax2, solution_path)
ax1.set(xlabel='x [Bohr]', ylabel='Energy [Hartree]', title=r'Potential, Eigenstate, $\langle x \rangle$',
xlim=x_lim, ylim=energy_lim)
ax2.set(xlabel='[Bohr]', title=r'$\sigma _{x}$', yticks=[], xlim=uncertainty_lim, ylim=energy_lim)
if export_pdf:
plt.savefig(f'{pdf_path}/schroedinger_solution.pdf', dpi=300)
if show_plot:
plt.show()
plt.close()
def main():
plot_solution('test', 'test', x_lim=(-5, 5), energy_lim=(0, 3.5), wavefunction_scale=1.0)
if __name__ == '__main__':
main()
|
