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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
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))
def plot_potential(ax: plt.Axes, solution_path: Path):
"""
Plots the potential line on to a given plt.Axes object from the solution file
:param ax: The axis on which the potential should be drawn
:param solution_path: The complete path where the solution is stored
"""
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: Path, wavefunction_scale: float = 1.0):
"""
Plots the wavefunc-data on a given plt.Axes object from the solution files
: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
"""
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)
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
def plot_expected_value(ax1: plt.Axes, ax2: plt.Axes, solution_path: Path):
"""
Plots the expected values on the given plt.Axes objects from the solution file.
:param ax1: Draws green 'x' symbols on an x-Energy plot
:param ax2: Draws yellow '+' symbols on an x-Energy plot
:param solution_path: The complete path where the solution is stored
:return:
"""
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_dir : str | None,
output_dir: str | None,
show_plot: bool,
export_pdf: bool,
wavefunction_scale: float,
energy_lim: None | tuple[float, float],
x_lim: None | tuple[float, float],
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.
: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)
"""
solution_path = Path(solution_dir) if solution_dir else Path.cwd()
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)
max_value, min_value = 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
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:
output_path = Path(output_dir) if output_dir else Path.cwd()
plt.savefig(output_path/'schroedinger_solution.pdf', dpi=300)
if show_plot:
plt.show()
plt.close()
def main():
# argparse setup
parser = argparse.ArgumentParser(description='Plots the solutions from schrodinger_solve.py')
msg = 'the path of the solution directory (default: None)'
parser.add_argument('-s', '--solution_dir', help=msg)
msg = 'the path where the pdf should be saved (default: None)'
parser.add_argument('-o', '--output_dir', 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()
# 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)
if __name__ == '__main__':
main()
|
