Writing the Schrodinger equation in the form: \begin{equation} \frac{\partial \psi}{\partial t} = i\mathcal{L}\psi+i\mathcal{N}\psi \end{equation} Where for the TDSE: \begin{equation} \mathcal{L} = \frac{1}{2m}\frac{\partial^2}{\partial x^2}, \ \ \ \ \ \mathcal{N} = -V(x) \end{equation}
By first neglecting $\mathcal{L}$ in time interval $[t_0,t_0+\Delta t/2]$ we are left with an ODE with a solution of the form:
\begin{equation} \psi (x,t_0 + \Delta t) = \exp(i\Delta t \mathcal{N}/2)\psi(x,t_0) \end{equation}Now neglecting $\mathcal{N}$, moving to momentum space $\mathcal{L}$ is simply multiplication. Hence in the full time interval $\Delta t$:
\begin{equation} \tilde{\psi} (k,t_0 + \Delta t) = \exp(i\Delta t \mathcal{F}(\mathcal{L}))\tilde{\psi}(k,t_0) = \exp(-i\Delta t k^2 /2m)\tilde{\psi}(k,t_0) \end{equation}For the initial $\tilde{\psi}(k,t_0)$ we use the Fourier transform of the time half step result we found first. Finally we must perform an additional spatial domain time half step to recover the split step approximation to time evolution by $\Delta t$.
In full, the process is the following:
\begin{equation} \psi (x,t_0+\Delta t) = \exp(i\Delta t \mathcal{N}/2) \mathcal{F}^{-1}(\exp(i\Delta t\mathcal{F}(\mathcal{L})) \ \mathcal{F} (\exp(i\Delta t \mathcal{N}/2) \ \psi(x,t_0))) \end{equation}We will be using Fast Fourier Transforms (FFTs) from the SciPy library so need to take into consideration the discrete nature of our input.
The basic argument behind this is to match the continuous Fourier transform pair $\psi(x,t) \leftrightarrow \tilde\psi(k,t)$ to a discrete approximation, $\psi(x_n,t) \leftrightarrow \tilde\psi(k_m,t)$. Here we use n and m to index x and k:
\begin{equation} \tilde\psi (k,t) = \frac{1}{\sqrt{2\pi}}\int^\infty_\infty \psi(x,t) e^{-ikx} dx \to \tilde\psi (k_m,t) \approx \frac{\Delta x}{\sqrt{2\pi}} \sum^{N-1}_{n=0} \psi(x_n,t) e^{-ik_mx_n} \end{equation}Comparing these to the discrete Fourier transform definitions we find the discrete Fourier transform pair:
\begin{equation} \frac{\Delta x}{\sqrt{2\pi}} \psi(x_n,t) e^{-ik_0x_n} \leftrightarrow \tilde\psi (k_m,t) e^{im\Delta k x_0} \end{equation}Where $\Delta k = 2\pi / (N \Delta x)$
#NAME: Schrödinger Equation
#DESCRIPTION: Solving the 1D time dependent Schrödinger Equation using the split step method.
import pycav.display as display
The top frame shows the spatial wavefunction, the real part in blue, imaginary in red and absolute value in black. Also showing is the potential shown in transparent black. In the bottom frame is the momentum wavefunction displayed in $k$ space
Gaussian wavepacket incident on potential barrier
display.display_animation('barrier.mp4')
Gaussian wavefunction placed at rest displaced from the centre of a harmonic potential
display.display_animation('linwell.mp4')
Starting with a Gaussian wavepacket incident on a square potential barrier, we expect some of the wavepacket to pass through and some to be reflected
import numpy as np
import pycav.pde as pde
import matplotlib.pyplot as plt
import matplotlib.animation as anim
def oneD_gaussian(x,mean,std,k0):
return np.exp(-((x-mean)**2)/(4*std**2)+ 1j*x*k0)/(2*np.pi*std**2)**0.25
def V(x):
V_x = np.zeros_like(x)
V_x[np.where(abs(x) < 0.5)] = 1.5
return V_x
We must define the spatial domain as well as the time step size and number of steps. Then the initial wavefunction is passed to the split step algorithm explained above.
N_x = 2**11
dx = 0.05
x = dx * (np.arange(N_x) - 0.5 * N_x)
dt = 0.01
N_t = 2000
p0 = 2.0
d = np.sqrt(N_t*dt/2.)
psi_0 = oneD_gaussian(x,x.max()-10.*d,d,-p0)
psi_x,psi_k,k = pde.split_step_schrodinger(psi_0, dx, dt, V, N_t, x_0 = x[0])
The time evolution can then be animated showing the expected behaviour. The momentum space wavefunction can also be visualised (here shown in the second subplot).
real_psi = np.real(psi_x)
imag_psi = np.imag(psi_x)
absl_psi = np.absolute(psi_x)
abs_psik = np.absolute(psi_k)
fig = plt.figure(figsize = (10,10))
ax1 = plt.subplot(211)
line1_R = ax1.plot(x,real_psi[:,0],'b')[0]
line1_I = ax1.plot(x,imag_psi[:,0],'r')[0]
line1_A = ax1.plot(x,absl_psi[:,0],'k')[0]
line_V = ax1.plot(x,0.5*V(x),'k',alpha=0.5)[0]
ax1.set_ylim((real_psi.min(),real_psi.max()))
ax1.set_xlim((x.min(),x.max()))
ax2 = plt.subplot(212)
line2 = ax2.plot(k,abs_psik[:,1],'k')[0]
ax2.set_ylim((abs_psik.min(),abs_psik.max()))
ax2.set_xlim((-10,10))
def nextframe(arg):
line1_R.set_data(x,real_psi[:,10*arg])
line1_I.set_data(x,imag_psi[:,10*arg])
line1_A.set_data(x,absl_psi[:,10*arg])
line2.set_data(k,abs_psik[:,10*arg])
animate = anim.FuncAnimation(fig, nextframe, frames = int(N_t/10), interval = 50, repeat = False)
animate = display.create_animation(animate,temp = True)
display.display_animation(animate)