-
Notifications
You must be signed in to change notification settings - Fork 0
/
wavefunction_plotter.py
65 lines (51 loc) · 1.69 KB
/
wavefunction_plotter.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
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from error import error
from probability_plotter import normalise
from analytical_psi import psi as analytical_psi
norm=True # should wavefunction be plotted as normalised?
comparison=True # display comparison plot?
"""
Plots the wavefunction data produced by "wavefunction_generator.py"
"""
try:
wave_function_csv="wavefunctions_N=1000.csv"
except:
error("Unable to load"+str(wave_function_csv))
data=pd.read_csv("wavefunctions\\"+str(wave_function_csv)) #create dataframe from csv
N=len(data)
x_array=np.linspace(-1/2,1/2,N) # positional array
plt.figure("Wavefunctions (N={})".format(N))
plt.xlabel(r"$\tilde{x}$")
plt.ylabel(r"$\psi(\tilde{x})$")
i=1
while 1:
try:
psi=data[str(i)]
norm_psi=psi*normalise(data[str(i)], x_array) #normalise wavefunction
if norm==True:
plt.plot(x_array,norm_psi, label="n="+str(i))
else:
plt.plot(x_array,psi, label="n="+str(i))
i+=1
except:
break
plt.tick_params(which='both',direction='in',right=True,top=True)
plt.legend()
if comparison==True:
"""
Plots difference between the numerically calculated wavefunction and it's analytical counterpart
"""
plt.figure("Wavefunctions comparison")
i=1
while i==1:
try:
psi_diff=np.array(data[str(i)]*normalise(data[str(i)], x_array) /analytical_psi(x_array, i, 1))
plt.plot(np.delete(x_array,[999]),np.log10(np.delete(psi_diff, [999])),label="n="+str(i))
i+=1
except:
break
plt.xlabel(r"$\tilde{x}$")
plt.ylabel(r"$log(\psi_{numerical}/\psi_{analytical})$")
plt.show()