{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Solving the 1D time-dependent Schrodinger equation using the split-step method\n",
"\n",
"Writing the Schrodinger equation in the form:\n",
"\\begin{equation}\n",
" \\frac{\\partial \\psi}{\\partial t} = i\\mathcal{L}\\psi+i\\mathcal{N}\\psi\n",
"\\end{equation}\n",
"Where for the TDSE:\n",
"\\begin{equation}\n",
" \\mathcal{L} = \\frac{1}{2m}\\frac{\\partial^2}{\\partial x^2}, \\ \\ \\ \\ \\ \\mathcal{N} = -V(x)\n",
"\\end{equation}\n",
"\n",
"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:\n",
"\n",
"\\begin{equation}\n",
" \\psi (x,t_0 + \\Delta t) = \\exp(i\\Delta t \\mathcal{N}/2)\\psi(x,t_0)\n",
"\\end{equation}\n",
"\n",
"Now neglecting $\\mathcal{N}$, moving to momentum space $\\mathcal{L}$ is simply multiplication. Hence in the full time interval $\\Delta t$:\n",
"\n",
"\\begin{equation}\n",
" \\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)\n",
"\\end{equation}\n",
"\n",
"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$.\n",
"\n",
"In full, the process is the following:\n",
"\n",
"\\begin{equation}\n",
" \\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)))\n",
"\\end{equation}\n",
"\n",
"We will be using Fast Fourier Transforms (FFTs) from the SciPy library so need to take into consideration the discrete nature of our input.\n",
"\n",
"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:\n",
"\n",
"\\begin{equation}\n",
" \\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}\n",
"\\end{equation}\n",
"\n",
"Comparing these to the discrete Fourier transform definitions we find the discrete Fourier transform pair:\n",
"\n",
"\\begin{equation}\n",
" \\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}\n",
"\\end{equation}\n",
"\n",
"Where $\\Delta k = 2\\pi / (N \\Delta x)$"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"#NAME: Schrödinger Equation\n",
"#DESCRIPTION: Solving the 1D time dependent Schrödinger Equation using the split step method.\n",
"\n",
"import pycav.display as display\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Example Animations\n",
"\n",
"### Plotting Information\n",
"\n",
"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\n",
"\n",
"Gaussian wavepacket incident on potential barrier"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"display.display_animation('barrier.mp4')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Gaussian wavefunction placed at rest displaced from the centre of a harmonic potential"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"display.display_animation('linwell.mp4')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Below is the code used to produce this animation\n",
"\n",
"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"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"import pycav.pde as pde\n",
"import matplotlib.pyplot as plt\n",
"import matplotlib.animation as anim\n",
"\n",
"def oneD_gaussian(x,mean,std,k0):\n",
" return np.exp(-((x-mean)**2)/(4*std**2)+ 1j*x*k0)/(2*np.pi*std**2)**0.25\n",
"\n",
"def V(x):\n",
" V_x = np.zeros_like(x)\n",
" V_x[np.where(abs(x) < 0.5)] = 1.5\n",
" return V_x"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"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."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"N_x = 2**11\n",
"dx = 0.05\n",
"x = dx * (np.arange(N_x) - 0.5 * N_x)\n",
"\n",
"dt = 0.01\n",
"N_t = 2000\n",
"\n",
"p0 = 2.0\n",
"d = np.sqrt(N_t*dt/2.)\n",
"\n",
"psi_0 = oneD_gaussian(x,x.max()-10.*d,d,-p0)\n",
"\n",
"psi_x,psi_k,k = pde.split_step_schrodinger(psi_0, dx, dt, V, N_t, x_0 = x[0])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"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)."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {
"collapsed": false
},
"outputs": [
{
"data": {
"text/html": [
""
],
"text/plain": [
""
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"real_psi = np.real(psi_x)\n",
"imag_psi = np.imag(psi_x)\n",
"absl_psi = np.absolute(psi_x)\n",
"abs_psik = np.absolute(psi_k)\n",
"\n",
"fig = plt.figure(figsize = (10,10))\n",
"ax1 = plt.subplot(211)\n",
"line1_R = ax1.plot(x,real_psi[:,0],'b')[0]\n",
"line1_I = ax1.plot(x,imag_psi[:,0],'r')[0]\n",
"line1_A = ax1.plot(x,absl_psi[:,0],'k')[0]\n",
"line_V = ax1.plot(x,0.5*V(x),'k',alpha=0.5)[0]\n",
"ax1.set_ylim((real_psi.min(),real_psi.max()))\n",
"ax1.set_xlim((x.min(),x.max()))\n",
"\n",
"ax2 = plt.subplot(212)\n",
"line2 = ax2.plot(k,abs_psik[:,1],'k')[0]\n",
"ax2.set_ylim((abs_psik.min(),abs_psik.max()))\n",
"ax2.set_xlim((-10,10))\n",
"\n",
"def nextframe(arg):\n",
" line1_R.set_data(x,real_psi[:,10*arg])\n",
" line1_I.set_data(x,imag_psi[:,10*arg])\n",
" line1_A.set_data(x,absl_psi[:,10*arg])\n",
" line2.set_data(k,abs_psik[:,10*arg])\n",
" \n",
"animate = anim.FuncAnimation(fig, nextframe, frames = int(N_t/10), interval = 50, repeat = False)\n",
"animate = display.create_animation(animate,temp = True)\n",
"display.display_animation(animate)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": []
}
],
"metadata": {
"anaconda-cloud": {},
"kernelspec": {
"display_name": "Python [default]",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.5.2"
}
},
"nbformat": 4,
"nbformat_minor": 0
}