-
Notifications
You must be signed in to change notification settings - Fork 0
/
solver.py
196 lines (172 loc) · 6.75 KB
/
solver.py
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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
import numpy as np, matplotlib.pyplot as plt, scipy.integrate as integrate
import warnings
import eel
warnings.filterwarnings('ignore')
class TISE:
def __init__(self, x_min, x_max, n_point, n_sol, e_min, e_max, de, const_b=2, potential=0) -> None:
self.x_min = x_min
self.x_max = x_max
self.n_point = n_point
self.n_sol = n_sol
self.const_b = const_b
self.x_array = np.linspace(self.x_min, self.x_max, num=self.n_point)
self.e_min = e_min
self.e_max = e_max
self.de = de
self.potential = potential
def get_pot(self, x):
if self.potential == 1:
return - np.exp(-x)
elif self.potential == 2:
return - 1 / (np.exp(x - 5) + 1)
elif self.potential == 3:
return - np.exp(x) / x
elif self.potential == 4:
return - 1 / (np.exp(x) - 1)
elif self.potential == 5:
return np.exp(-2 * (x - 1)) - 2 * np.exp(-(x - 1))
elif self.potential == 6:
return (1 - 2 * x) / x ** 2
elif self.potential == 7:
return - 1 / x + x
elif self.potential == 8:
return x ** (-12) - 2 * x ** (-6)
elif self.potential == 9:
return - 1 / x
else:
return x ** 2 / 2
def solve(self):
numerov = Numerov(x_min=self.x_min, x_max=self.x_max, e_min=self.e_min, e_max=self.e_max, de=self.de, const_b=self.const_b, n_sol=self.n_sol, potential=self.potential)
try:
self.x_array, self.psi, self.eigs = numerov.get_sol()
self.normalize()
return len(self.eigs)
except:
return 0
def normalize(self):
norm_psi = list()
for eig_func in self.psi:
integral = integrate.simpson(np.abs(eig_func)**2, self.x_array)
norm_psi.append(eig_func / integral)
self.psi = norm_psi
def plot_pot(self):
plt.figure(figsize=(16, 9))
plt.plot(self.x_array, self.get_pot(self.x_array), ls='-', lw=2, c='red')
plt.grid()
plt.xlabel('x')
plt.ylabel('F(x)')
plt.show()
def plot_solutions(self):
f = self.get_pot(self.x_array)
fig, ax = plt.subplots(figsize=(12,6))
ax.plot(self.x_array, f, color="black", lw=4, ls='--')
ax.set_xlabel("x", fontsize=14)
ax.set_ylabel("F(x)", fontsize=14)
ax.grid()
ax2 = ax.twinx()
try:
for ind, sol in enumerate(self.psi):
ax2.plot(self.x_array, sol, lw=4, label=f'E={self.eigs[ind]}')
ax2.set_ylabel("$\Psi(x)$", fontsize=14)
ax2.legend()
plt.show()
except:
eel.show_message(0)
# The Numerov class implements algorithm presented in the paper (check link in the top left of the gui)
class Numerov():
def __init__(self, x_min, x_max, e_min, e_max, de, n_point=1024, const_b=2, n_sol=4, precision=1e-6, potential=0):
self.x_min = x_min
self.x_max = x_max
self.n_point = n_point
self.const_b =const_b
self.n_sol = n_sol
self.x_array = np.linspace(x_min, x_max, n_point)
self.dx = self.x_array[1] - self.x_array[0]
self.precision = precision
self.potential = potential
self.e_min = e_min
self.e_max = e_max
self.de = de
self.pot = TISE(x_min=x_min, x_max=x_max, n_point=n_point, n_sol=n_sol, e_min=self.e_min, e_max=self.e_max, de=self.de, const_b=const_b, potential=self.potential).get_pot(self.x_array)
self.eps = e_min
try:
self.c = np.where(self.pot >= self.eps)[0][0]
if self.c < 2:
self.c = 2
elif self.c > self.n_point - 2:
self.c = self.n_point - 2
except:
self.c = self.n_point - 2
self.small_a = 1e-5
def factor(self, ind):
return self.const_b * (self.eps - self.pot[ind])
# Forward relation (check link in the top left of the gui)
def forward_relation(self):
y = np.zeros(self.c)
y[0] = 0
y[1] = self.small_a
for i in range(2, self.c):
cf = (1 + self.dx ** 2 / 12 * self.factor(i))
y[i] = 2 * (1 - 5 * self.dx ** 2 / 12 * self.factor(i-1)) * y[i-1] - (1 + self.dx ** 2 / 12 * self.factor(i-2)) * y[i-2]
y[i] = y[i] / cf
gl = (y[-1] - y[-2]) / (self.dx * y[-1])
return y, gl
# Backward relation (check link in the top left of the gui)
def backward_relation(self):
y = np.zeros(self.n_point - self.c + 1)
y[-1] = 0
y[-2] = self.small_a
for j in range(3, self.n_point - self.c + 2):
i = -j
cf = (1 + self.dx ** 2 / 12 * self.factor(i))
y[i] = 2 * (1 - 5 * self.dx ** 2 / 12 * self.factor(i+1)) * y[i+1] - (1 + self.dx ** 2 / 12 * self.factor(i+2)) * y[i+2]
y[i] = y[i] / cf
y = y / np.max(y)
gr = (y[1] - y[0]) / (self.dx * y[1])
return y, gr
# Algorithm according to Numerov
def numerov_algorithm(self):
try:
self.c = np.where(self.pot >= self.eps)[0][0]
except:
self.c = self.n_point - 2
if self.c < 2:
self.c = 2
yl, gl = self.forward_relation()
yr, gr = self.backward_relation()
yr = yr * yl[-1]
err = gl - gr
return np.concatenate([yl, yr[1:]]), err
# My custom function
def find_eigenvalues(self):
de0 = self.de
self.eps = self.e_min
loss_this = self.numerov_algorithm()[1]
e_eig = []
while self.eps < self.e_max and len(e_eig) < self.n_sol:
eel.setProgress(str(100 - int((self.e_max - self.eps) / (self.e_max - self.e_min) * 100)))
de = de0
loss_prev = loss_this
self.eps += de
loss_this = self.numerov_algorithm()[1]
if np.sign(loss_prev) != np.sign(loss_this):
de = -de / 2
while np.abs(de) > self.precision:
self.eps += de
loss_prev, loss_this = loss_this, self.numerov_algorithm()[1]
if np.sign(loss_prev) != np.sign(loss_this):
de = -de / 2
if np.abs(np.abs(self.backward_relation()[1]) - np.abs(self.forward_relation()[1])) < 1.0:
e_eig.append(self.eps)
de = de0
self.eps += de
loss_this = self.numerov_algorithm()[1]
eel.setProgress(str(100))
return e_eig
def get_sol(self):
eigs = self.find_eigenvalues()
psi = list()
for eig in eigs:
self.eps = eig
psi.append(self.numerov_algorithm()[0])
return self.x_array, psi, eigs