A Simulation of Solutions to the Schrodinger Equation¶
This project is meant to be a useful aid in visualizing the solutions of the time-dependant Schrodinger equation in basic cases. It is also meant to showcase the mathematics behind solving the equation using the Crank-Nicolson finite-difference method from the perspective of a first-year Physics undergraduate and amateur scientific programmer.
The equation in the 1-D case: $$\frac{-\hbar ^2}{2m} \frac{\partial ^2}{\partial x^2} \Psi +U(x) \Psi =i\hbar \frac{\partial}{\partial t}\Psi$$
This is a second-degree PDE, which is difficult to solve analytically. Instead, we will be solving this numerically. The Crank-Nicolson method is a method to solve such PDEs numerically, which we will be using to simulate a variety of simple 1D problems (and one 3D problem).
The Crank-Nicolson Method¶
In our case of $\Psi (x,t)$, the idea is to numerically solve the equation (finding the numerical value of $\Psi$ for some given $x$ and $t$). Then, assuming $\Psi$ is smooth, we interpolate between these known points.
Let's say our domain of interest is rectangular, $x \in [x_0,x_0+w]$, and $t \in [0, t_f]$. We subdivide these domains into $S_x$ and $S_t$ equal subintervals of length $L_x$ and $L_t$ respectively, and designate $\Psi_{(j,k)}$ as the value of $\Psi$ at $(x,t) = (x + j L_x, k L_t)$. That is, $(j,k)$ is the index on the "grid" created by subdividing the domain:
We can then create expressions to approximate the relevant partial derivatives of $\Psi$ at every $(j,k)$. Note that we actually center the difference expressions around the midpoints of each $t$-interval:
If we plug these back into our equation, then move the $k$-terms to one side and the $k+1$ terms to the other, we can re-arrange it as a linear system $Ax=B$ where $B$ has entries of the side of the equation dependant on $\Psi (\cdot,k)$, $x$ is a vector with entries $\Psi(\cdot, k+1)$ and $A$ has the coefficients of $\Psi(\cdot, k+1), \Psi(\cdot, k)$, and $\Psi(\cdot, k-1.)$ Then, we set appropriate boundary conditions, and can solve each successive time state based on the state of the previous state.
First we simulate the particle in a box. This simply has $U(x,t) = 0, \Psi = 0$ at $x = 0$ and $x = L$.
The Particle In A Box¶
We start with our equation: $$\frac{-\hbar ^2}{2m} \frac{\partial ^2}{\partial x^2} \Psi =i\hbar \frac{\partial}{\partial t}\Psi$$ Plugging in our derivatives: $$\frac{-\hbar ^2}{2m} \frac{\Psi_{(j + 1, k+1)} - 2\Psi_{(j, k+1)} + \Psi_{(j-1,k+1)} + \Psi_{(j + 1, k)} - 2\Psi_{(j, k)} + \Psi_{(j-1,k)}}{2(L_x)^2} =i\hbar \frac{\Psi_{(j,k+1)} - \Psi_{(j,k)}}{L_t}$$ Moving the $k$ terms to the left and $k+1$ to the right: $$\frac{-\hbar ^2}{2m} \frac{\Psi_{(j + 1, k)} - 2\Psi_{(j, k)} + \Psi_{(j-1,k)}}{2(L_x)^2} + i\hbar \frac{\Psi_{(j,k)}}{L_t} = \frac{\hbar ^2}{2m} \frac{\Psi_{(j + 1, k+1)} - 2\Psi_{(j, k+1)} + \Psi_{(j-1,k+1)} }{2(L_x)^2} + i\hbar \frac{\Psi_{(j,k+1)}}{L_t}$$ Grouping each side by $j$, $j+1$, and $j-1$:
$$\frac{-\hbar ^2}{2m} \frac{\Psi_{(j + 1, k)}}{2(L_x)^2} + \frac{-\hbar ^2}{2m} \frac{\Psi_{(j - 1, k)}}{2(L_x)^2} + (\frac{i\hbar}{L_t} + \frac{2\hbar^2}{2m*2(L_x)^2})\Psi_{(j,k)} = \frac{\hbar ^2}{2m} \frac{\Psi_{(j + 1, k+1)}}{2(L_x)^2} + \frac{\hbar ^2}{2m} \frac{\Psi_{(j - 1, k+1)}}{2(L_x)^2} + (\frac{i\hbar}{L_t} + \frac{-2\hbar^2}{2m*2(L_x)^2})\Psi_{(j,k+1)}$$
Now multiplying by $4mL_x^2L_t$: $$-\hbar^2L_t\Psi_{(j + 1, k)} + -\hbar^2L_t\Psi_{(j - 1, k)} + (4i\hbar mL_x^2 + 2\hbar^2L_t)\Psi_{(j,k)} = \hbar^2L_t\Psi_{(j + 1, k+1)/} + \hbar^2L_t\Psi_{(j - 1, k+1)} + (4i\hbar mL_x^2 - 2\hbar^2L_t)\Psi_{(j,k+1)}$$
Let's designate the entire left side as $D_{(j,k)}$. Then on the right, let $$A_j = \hbar^2L_t$$ $$B_j = (4i\hbar mL_x^2 - 2\hbar^2L_t)$$ $$C_j = \hbar^2L_t$$
For some value of $D_j$ at time stepping $k$, it is equal to $A_j\Psi_{(j-1,k+1)} + B_j\Psi_{(j,k+1)} + C_j\Psi_{(j+1,k+1)}$
We can then represent this as a linear system and solve for each $\Psi_{(\cdot,k+1)}$ based on known $D_{(\cdot,k)}$.
Our stationary state is given by $\Psi = \sqrt{\frac{2}{L}}\sin{(\frac{\pi nx}{L})}$
from matplotlib import pyplot as plt, animation as anim
import numpy as np
import ipywidgets as widgets
from IPython.display import display
#CONSTANTS
#I have no clue what the units are any more? time I set to nanoseconds, length to nm.
m = 0.5 #I guess these are something?
HBAR = 6.582119569E-7#6.582119569E-6 #this is 1e10 greater than its actual value in eV/s, but I need to scale everything up so we stay within FP precision
l = 1 #why did i bother setting this? This is in nm.
wl = 0.2 #the length of the well, it's centered at x = 0.5
U = 0 #this is potential, probably needs to be turned into a matrix
Lt = 20 #this is the stepping for time. What units? no clue. At least now not horrifically huge.
Sx = 200 #number of steppings in x
Lx = l/Sx #this is the granularity of our finite-difference stepping in x (step-size)
#coefficients of the matrix
Ai = (np.power(HBAR,2)) * Lt
Bi = (-2*np.power(HBAR,2)*Lt) + (4*np.power(Lx,2)*1j*HBAR*m)
Bl = (2*np.power(HBAR,2)*Lt) + (4*np.power(Lx,2)*1j*HBAR*m) #this one is for Di
Ci = (np.power(HBAR,2)) * Lt
n = 2 #just for the stationary state for now...
psi_type = 0 #to swap between psi types
cached_psi = 0 #because i don't know how else to do real-time checking
u_type = 0
cached_u = 0
times_crunched = 0
SF = np.longdouble((np.power(10,8)))
#INITIALIZATION
#this line fixes everything. "Every[thing]? EVERY[thing]!!!!"
%matplotlib notebook
cs = []
da = []
ta = []
ua = []
#uhh how do i make a gui
def init_psi(type):
global CS, DA
for i in range (0, Sx):
if type == 0:
cs.append((np.sqrt(2/l) * np.sin(n*np.pi*(Lx*i)/l)))
elif type == 1:
cs.append((np.sqrt(2/l) * np.exp(-np.pow(((Lx*i)-0.11)/(0.13), 2))))
elif type == 2:
cs.append((np.sqrt(2/l) * np.exp(-100*np.pow(Lx*i - 0.5,2))))
da.append(-1)
init_psi(0)
CS = np.array(cs) #holds all the psi values currently known
DA = np.array(da, complex) #holds the rhs of the equation
def reset_psi(type):
global CS, DA
for i in range (0, Sx):
if type == 0:
CS[i]=(np.sqrt(2 / l) * np.sin(n * np.pi * (Lx * i) / l))
elif type == 1:
CS[i]=(np.sqrt(2 / l) * np.exp(-np.pow(((Lx * i) - 0.11) / (0.13), 2)))
elif type == 2:
CS[i]=(np.sqrt(2 / l) * np.exp(-100 * np.pow(Lx * i - 0.5, 2)))
DA[i] = -1
def init_u(type):
global ua, l
if type == 0:
for i in range (0, Sx):
ua.append(0)
elif type == 1:
for i in range (0, Sx):
cx = (Lx * i) - (l/2)
cf = np.pow(cx,2)
ua.append(4*cf/SF)
elif type == 2:
for i in range (0, Sx):
if (i*Lx != 0.5-(wl/2)) & (i*Lx != 0.5+(wl/2)):
ua.append(0)
else:
ua.append(1/SF)
init_u(0)
UA = np.array(ua, np.longdouble)
def reset_u(type):
global UA, Sx, Lx, l
if type == 0:
for i in range (0, Sx):
UA[i] = 0
elif type == 1:
for j in range (0, Sx):
cx = (Lx * j) - (l/2)
cf = np.pow(cx,2)
UA[j]=(4*cf)/SF
elif type == 2:
for k in range (0, Sx):
if (k*Lx < 0.5-(wl/2)) | (k*Lx > 0.5+(wl/2)):
UA[k] = 0
else:
UA[k] = 1/SF
def init_matrix():
global TA
for i in range (0, Sx):
temp = []
for j in range (0, Sx):
match i-j:
case -1:
temp.append(Ci)
case 0:
temp.append(Bi - (2*m*(np.power(Lx,2))*Lt*UA[j]))
case 1:
temp.append(Ai)
case _:
temp.append(0)
ta.append(temp)
#boundary conditions
ta[0][0] = 1
ta[0][1] = 0
ta[Sx-1][Sx-2] = 0
ta[Sx-1][Sx-1] = 1
init_matrix()
TA = np.array(ta) #linear transformation to go from each state
def reset_matrix():
global TA
for i in range (0, Sx):
temp = []
for j in range (0, Sx):
match i-j:
case -1:
temp.append(Ci)
case 0:
temp.append(Bi - (2*m*(np.power(Lx,2))*Lt*UA[j]))
case 1:
temp.append(Ai)
case _:
temp.append(0)
TA[i]= temp
#boundary conditions
TA[0][0] = 1
TA[0][1] = 0
TA[Sx-1][Sx-2] = 0
TA[Sx-1][Sx-1] = 1
#ITERATION
#wow I suck at this whole python thing haha it literally took one line to make animations work
fig, ax = plt.subplots()
x = np.arange(0, Sx) #an array of 0, 1, 2, ... Sx
line, = ax.plot((x*l/Sx), CS)
line2, = ax.plot((x*l/Sx),UA)
psi_select = widgets.RadioButtons(options=[('Sinusoidal Eigenstate', 0), ('Gaussian',1), ('Centered Gaussian', 2)],description='Wave function:', disabled=False, continuous=True,)
U_select = widgets.RadioButtons(options=[('Particle in a Box', 0), ('Harmonic [BROKEN]',1), ('Finite Square Barrier(s)', 2)],description='Potential Function:', disabled=False, continuous=True,)
display(psi_select)
display(U_select)
pc = ax.annotate('',xy=(1, 0),
xycoords='axes fraction',
xytext=(-20, 420), textcoords='offset pixels',
horizontalalignment='right',
verticalalignment='bottom')
ss = ax.annotate('',xy=(1, 0),
xycoords='axes fraction',
xytext=(-20, 460), textcoords='offset pixels',
horizontalalignment='right',
verticalalignment='bottom')
def update(i):
global times_crunched, pc
times_crunched = times_crunched + 1
pc.set_text('iterations:'+str(times_crunched))
global psi_type, cached_psi, CS, u_type, cached_u
#for now we have to check every update cycle whether or not to reset. Inefficient, but it works.
psi_type = psi_select.value
u_type = U_select.value
if (psi_type != cached_psi) | (u_type != cached_u):
cached_psi = psi_type
reset_psi(psi_type)
cached_u = u_type
reset_u(u_type)
reset_matrix()
times_crunched = 0
#updates the right side of the linear system (RHS of augmented matrix)
for k in range(0, 1): #change the range to increase the time covered, but decrease the time granularity
DA[0] = 0
DA[Sx-1] = 0
for i in range (1,Sx-1):
DA[i] = (-Ai * CS[i-1]) + ((Bl+ (2*m*(np.power(Lx,2))*Lt*UA[i])) * CS[i]) - (Ci * CS[i+1])
CS = np.linalg.solve(TA,DA) #perhaps I will implement this myself at some point, but for now I will just call a library method.
psi_arr = np.abs(CS)
psi_squared = 0
for a in psi_arr:
psi_squared = psi_squared + ((a**2)*Lx)
ss.set_text('$|\\Psi|^{2}$ is '+str(psi_squared))
line.set_ydata(psi_arr)
return line,
def update_pot(i):
global UA
line2.set_ydata(UA * SF)
return line2,
plt.xlabel('x')
plt.ylabel('$|\\Psi|$')
ax.set_xlim(0,l)
ax.set_ylim(0,2)
angy2 = anim.FuncAnimation(fig, update, frames = np.linspace(1, 1000, 1000), interval = 25)
angy3 = anim.FuncAnimation(fig, update_pot, frames = np.linspace(1, 1000, 1000), interval = 25)
plt.show()
If you try setting it to a pure energy eigenstate such as $\sin{(\frac{n\pi x}{L})}, n \in \mathbb{Z}$ you can see the wave function basically doesn't change (what change is present is of course caused by my finite difference method). Alternatively, you can set a Gaussian wave packet and watch it evolve over time.
More Boundaries¶
Implementing more boundary conditions involves simply modifying the boundary conditions. However, due to the nature of computing within a finite domain, there will still be a hard boundary at the edges, which can cause reflection artifacts. To rectify this, a Complex Absorbing Potential (CAP) is needed. My math skillz suck, so this won't be done for a bit. Instead, time was spent on fancy labels for the simulation.
The General 1D Solution¶
Now, we include $U(x)$: $$\frac{-\hbar ^2}{2m} \frac{\partial ^2}{\partial x^2} \Psi + U(x)\Psi =i\hbar \frac{\partial}{\partial t}\Psi$$ Plugging in our derivatives (and now $\Psi$): $$\frac{-\hbar ^2}{2m} \frac{\Psi_{(j + 1, k+1)} - 2\Psi_{(j, k+1)} + \Psi_{(j-1,k+1)} + \Psi_{(j + 1, k)} - 2\Psi_{(j, k)} + \Psi_{(j-1,k)}}{2(L_x)^2} + U(x)\frac{\Psi_{(j,k)} + \Psi_{(j,k+1)}}{2}=i\hbar \frac{\Psi_{(j,k+1)} - \Psi_{(j,k)}}{L_t}$$ Moving the $k$ terms to the left and $k+1$ to the right: $$\frac{-\hbar ^2}{2m} \frac{\Psi_{(j + 1, k)} - 2\Psi_{(j, k)} + \Psi_{(j-1,k)}}{2(L_x)^2} + i\hbar \frac{\Psi_{(j,k)}}{L_t} + U(x)\frac{\Psi_{(j,k)}}{2} = \frac{\hbar ^2}{2m} \frac{\Psi_{(j + 1, k+1)} - 2\Psi_{(j, k+1)} + \Psi_{(j-1,k+1)} }{2(L_x)^2} + i\hbar \frac{\Psi_{(j,k+1)}}{L_t} - U(x)\frac{\Psi_{(j,k+1)}}{2}$$ Grouping each side by $j$, $j+1$, and $j-1$:
$$\frac{-\hbar ^2}{2m} \frac{\Psi_{(j + 1, k)}}{2(L_x)^2} + \frac{-\hbar ^2}{2m} \frac{\Psi_{(j - 1, k)}}{2(L_x)^2} + (\frac{i\hbar}{L_t} + \frac{2\hbar^2}{2m*2(L_x)^2} + \frac{U(x)}{2})\Psi_{(j,k)} = \frac{\hbar ^2}{2m} \frac{\Psi_{(j + 1, k+1)}}{2(L_x)^2} + \frac{\hbar ^2}{2m} \frac{\Psi_{(j - 1, k+1)}}{2(L_x)^2} + (\frac{i\hbar}{L_t} + \frac{-2\hbar^2}{2m*2(L_x)^2} - \frac{U(x)}{2})\Psi_{(j,k+1)}$$
Now multiplying by $4mL_x^2L_t$: $$-\hbar^2L_t\Psi_{(j + 1, k)} + -\hbar^2L_t\Psi_{(j - 1, k)} + (4i\hbar mL_x^2 + 2\hbar^2L_t + 2mL_x^2L_tU(x))\Psi_{(j,k)}$$ $$= \hbar^2L_t\Psi_{(j + 1, k+1)/} + \hbar^2L_t\Psi_{(j - 1, k+1)} + (4i\hbar mL_x^2 - 2\hbar^2L_t - 2mL_x^2L_tU(x))\Psi_{(j,k+1)}$$
Again, the entire left side is $D_{(j,k)}$. Then on the right, let $$A_j = \hbar^2L_t$$ $$B_j = (4i\hbar mL_x^2 - 2\hbar^2L_t - 2mL_x^2L_tU(x))$$ $$C_j = \hbar^2L_t$$
And we do the same as before.