-
Notifications
You must be signed in to change notification settings - Fork 0
/
wavefunction_generator.py
189 lines (152 loc) · 5.51 KB
/
wavefunction_generator.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
from potential_class import potential
from quantum_particle import q_particle
import numpy as np
import copy
import pandas as pd
from error import error
import time
well_length=1 #length of the potential well
E_list=[0.9,0.8,0.5,0.2,-0.2,-0.7,-1.4,-2.1,-2.9,-3.9] # list of energy level starting points
psi_tolerance=1E-12 # tolerance on boundary condition ψ(L/2)=0
psi_dict={}
if __name__=="__main__":
"""
Setup if script is being run directly
"""
start=time.time() #begins timing
N=round(1000) #number of integration steps
x_array=np.linspace(-well_length/2,well_length/2,N) #array of positions
V=potential(N) #potential object
def turn_points(array):
"""
Counts the number of turning points in an array
Parameters
----------
array : numpy.ndarray
Array of values to be processed.
Returns
-------
n : int
Number of turning points.
"""
n=0
for i in range(1,len(array)-1):
if (array[i]<array[i-1] and array[i]<array[i+1]) or (array[i]>array[i-1] and array[i]>array[i+1]):
n+=1
return n
def WF_attempt(trial_E):
"""
Creates a Quantum Harmonic Oscillator object and attempts to
produce a valid wavefunction using Numerov's formula
Parameters
----------
trial_E : float, optional
The dimentionless energy for this attempt.
Returns
-------
wave: numpy.ndarray
Final wavefunction attempt.
"""
electron=q_particle(trial_E,V.V_depth(),well_length,N) #create particle object
for i in range(2,N): # attempt to generate valid wavefunction with Numerov algorithm
electron.next_psi(V.nu(), i)
wave=copy.deepcopy(electron.wavefunction)
del electron # delete object to save on memory during large operations
return wave
def E_finder(inital_E):
"""
Iterates WF_attempt with changing energy until boundary condition ψ(L/2)=0 is met within tolerance
Parameters
----------
inital_E : float
DESCRIPTION.
Returns
-------
trial_E : float
The energy eigenvalue of valid wavefuction.
test_wave : np.ndarray
Valid, non-normalised wavefunction.
n : int
Principal quantum number of wavefunction .
"""
trial_E=inital_E
delta_E=trial_E/100
best_psi=WF_attempt(trial_E)[-1] # value of the best fulfilment of the ψ(L/2)=0 boundary condition
last_psi=copy.deepcopy(best_psi) # ψ(L/2) for last attempt
while abs(best_psi)>psi_tolerance: #loop until boundary condition is met within tolerance
trial_E+=delta_E
test_wave=WF_attempt(trial_E)
test_wave_L=test_wave[-1]
"""
Guides energy level closer to meeting the ψ(L/2)=0 boundary condition
"""
if abs(test_wave_L)>abs(best_psi):
delta_E=-delta_E
# print("Worse - Reversing direction")
else:
# print("Improvement")
best_psi=test_wave_L
if np.sign(test_wave_L)!=np.sign(last_psi):
delta_E=delta_E/2
# print("Overstep")
last_psi=copy.deepcopy(test_wave_L)
else:
n=turn_points(test_wave)
print("N={} - Value found (n={}): \nE = {}\nψ(L) = {}\n".format(N,n,trial_E,best_psi)) # ADD IN TURNING POINTS AND DATAFRAME
energy_dict["n"].append(n)
energy_dict["epsilon"].append(trial_E) #stores valid wavefunction and energy levels
psi_dict[str(n)]=test_wave
return trial_E,test_wave,n
def run(i):
"""
Finds wavefunctions of different quantum numbers.
Can be called as child for multiprocessing
Parameters
----------
i : int
No. of intergration steps/50. Used to indicated process no. when multiprocessing.
Returns
-------
N : int
No. of intergration steps.
time_elapsed : float
Time taken to identify all valid wavefunctions in required range at a given N value.
"""
global N
global x_array
global V
global energy_dict
start=time.time()
N=round(i*50)
x_array=np.linspace(-well_length/2,well_length/2,N)
V=potential(N)
energy_dict={"n":[], "epsilon":[]}
print("running N={}".format(N))
for energy in E_list:
E,psi,n=E_finder(energy)
energy_df=pd.DataFrame(data=energy_dict)
psi_df=pd.DataFrame(data=psi_dict)
row_limit=len(E_list)
if energy_df.shape[0]>row_limit: #catches extra iterations
rfd=[i for i in range(row_limit,energy_df.shape[0])] #Rows For Deletion
energy_df.drop(rfd,inplace=True)
energy_df.index.name=str(N)
try:
pd.DataFrame.to_csv(energy_df, "energy_levels\\energy_levels_N={}.csv".format(N)) # saves data in relavent csvs for storage
print("Saved \"energy_levels_N={}.csv\"".format(N))
except:
print("energy_levels\\energy_levels_N={}.csv".format(N))
error("Unable to save \"energy_levels_N={}.csv\"".format(N),False)
del energy_df
try:
pd.DataFrame.to_csv(psi_df, "wavefunctions\\wavefunctions_N={}.csv".format(N))
print("Saved \"wavefunctions_N={}.csv\"".format(N))
except:
error("Unable to save \"wavefunctions_N={}.csv\"".format(N),False)
del psi_df
time_elapsed=time.time()-start
return N,time_elapsed
if __name__=="__main__": #does not trigger if process is controlled by multiprocessing parent script
run(N/50)
print("Time elapsed = {}s".format(round(time.time()-start,2))) #finishes process timing
print("\a")