-
Notifications
You must be signed in to change notification settings - Fork 0
/
Lab09_Q2.py
110 lines (97 loc) · 3.1 KB
/
Lab09_Q2.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
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 14 09:46:30 2018
@author: brayd
"""
import numpy as np
from scipy.constants import m_e, hbar
#instantiate some values of constants
L = 1e-8 #length of box in meters
N = 1000 #spatial slices
a = L/N #size of grid boxes
M = m_e #mass = mass of electron in kg
#instantiate emtpy vector for psi(x,0)
psi = np.zeros(N+1, dtype = complex)
#define the function that psi(x,0) is represented by to later fill up the vector
def fun_psi_0(x):
'''
Calculate the wavefunction at t=0, for a value of x.
x should be a real valued function between zero and L for this to make
sense, physically.
'''
x_o = L/2 #mean of Guassian part of the wavefunction in meters
sigma = 1e-10 #stdev of Gaussian part of the wavefunction in meters
k = 5e10 #m^{-1}
out = np.exp(-(x - x_o)**2/(2*sigma**2))*np.exp(1j*k*x)
#NB: we forget about a normalization constant because it doesn't affect the solution
return out
#fill up the array, leaving the boundary conidtions in place
x = np.linspace(0, L, N+1)
psi[:] = fun_psi_0(x)
psi[0] = 0
psi[N] = 0
#normalize
norm = np.trapz(abs(psi)**2)
psi = psi/norm
#define v := B*psi(x,0), B as given in Newman pg.440
h = 1e-18 #time step in seconds
b1 = 1 - h*1j*hbar/(2*M*a**2)
b2 = h*1j*hbar/(4*M*a**2)
v = np.zeros(N, dtype = complex)
#calculate v = B.psi(x,0)
v = b1*psi[1:N] + b2*(psi[2:N+1] + psi[0:N-1])
#now to solve the system Ax = v
#x will be time evolved value of psi, A is as given in Newman pg.440
from banded import banded
a1 = 1 + h*1j*hbar/(2*M*a**2)
a2 = -h*1j*hbar/(4*M*a**2)
A = np.empty([3,N], dtype = complex)
A[0,:] = a2
A[1,:] = a1
A[2,:] = a2
psi_t = np.zeros(N+1, dtype = complex) #psi at later time t
#Calculate 1 time step:
psi[1:N] = banded(A, v, 1, 1) #solve for x using Gauss elim and backsub for tridiagnonal matrix
#the above psi_t is the wavefunction on the domain x in [0,L] for t=h (h time step)
#Extend program to perform repeated steps and make an animation
steps = 2000
from pylab import figure, plot, xlabel, ylabel, title, clf, pause, draw
figure()
plot(x, psi.real)
xlabel('$x$ m')
ylabel('$\psi(x,t)$')
title('Wavefunction at time t = 0')
#instantiate emtpy vector for psi(x,t)
psi_t = np.zeros(N+1, dtype = complex)
#fill up the array, leaving the boundary conidtions in place
psi_t[:] = fun_psi_0(x)
psi_t[0] = 0
psi_t[N] = 0
#normalize
norm = np.trapz(abs(psi_t)**2)
psi_t = psi_t/norm
#figure out the indeces at which we get to 1e-15s and 1e-16s to pause and save fig
time1 = 1e-15
time2 = 1e-16
ind1 = time1/h
ind1 = int(ind1)
ind2 = time2/h
ind2 = int(ind2)
figure()
for i in range(steps):
#calculate v = B.psi(x,t)
v = b1*psi_t[1:N] + b2*(psi_t[2:N+1] + psi_t[0:N-1])
#solve for x using Gauss elim and backsub for tridiagnonal matrix
psi_t[1:N] = banded(A, v, 1, 1)
the_time = (i+1)*h
#create animation
clf()
plot(x,psi_t.real)
xlabel('$x$ meters')
ylabel('$\psi(x,t)$')
title('Wavefunction at time t = {0:.2e}'.format(the_time))
draw()
if i == ind1 or i == ind2:
pause(10) #allow me to take screen shot
else:
pause(0.01)