Schrodinger equation in 1D

$$ i\hbar\frac{\partial\psi}{\partial t} = \hat H\psi $$

In particular $$ i\hbar\frac{\partial\psi}{\partial t} = \left[-\frac{\hbar^2}{2m}\Delta + V\right]\psi $$

In [17]:
%matplotlib inline
from pylab import *
from numpy import *
In [18]:
#physical constants
hbar = 1.0
m = 1.0
In [19]:
#discretization parameters
nx = 100   # number of unknown grid points in spatial direction
SC = 1.0  # stability criterion SC = 2*D*dt/dx^2, for FTCS should be SC <= 1

def schr_init(maxx, maxt, maxV, nx, SC):
    #choose time step according to CFL condition
    dx = maxx/(nx+1)
    dt = hbar/(2*hbar**2/(m*dx**2)+maxV)*SC
    #dt = SC*dx**2*m/hbar  # for V = 0
    nt = int(maxt/dt)+1
    
    #define array for storing the solution
    Ur = zeros((nt, nx+2))
    Ui = zeros((nt, nx+2))
    
    x = arange(nx+2)*dx
    t = arange(nt)*dt
    return Ur, Ui, dx, dt, x, t
In [20]:
def schr_solve(Ur, Ui, dx, dt, nx, nt):
    xint = arange(1, nx+1)
    for it in range(0,nt-1):
        Ui[it+1,xint] = Ui[it,xint] + hbar*dt/(2*m*dx**2)*(Ur[it,xint+1] - 2*Ur[it, xint] + Ur[it, xint-1])\
            - dt/hbar*V[xint]*Ur[it,xint]
        Ur[it+1,xint] = Ur[it,xint] - hbar*dt/(2*m*dx**2)*(Ui[it+1,xint+1] - 2*Ui[it+1, xint] + Ui[it+1, xint-1])\
            + dt/hbar*V[xint]*Ui[it+1,xint]

Assume the initial wave packet will be in the form $$ \psi(x, t_0) = \frac{1}{\sqrt{\sigma\sqrt{\pi}}}e^{-(x-x_0)^2/2\sigma}e^{ik_0x} $$

In [21]:
maxx = 2.
maxt = 0.05
nx = 200
Ur, Ui, dx, dt, x, t = schr_init(maxx, maxt, 0, nx, 1.0)
V = zeros_like(x)
In [22]:
sigma = 0.1
k0 = 50.
x0 = maxx/3.
U0 = 1./sqrt(sigma*sqrt(pi))*exp(-(x-x0)**2/(2*sigma**2)) * exp(1j*k0*x)
def Ut(t, x, x0=x0, k0=k0):
    p0 = hbar*k0
    norm = sqrt(sigma/(sqrt(pi)*(sigma**2+1j*hbar*t/m)))
    envelope = exp(-0.5*(x-x0-p0*t/m)**2/(sigma**2+1j*hbar*t/m))
    #phase = exp(1j/hbar*(p0*x-p0**2*t)/(2*m))
    phase = exp(1j*k0*(x-hbar*k0*t/(2*m)))
    return norm*envelope*phase
In [23]:
#plot(U0.real)
#plot(U0.imag)
plot(Ut(0, x).imag)
plot(Ut(0, x).real)
plot(abs(Ut(0, x)))
Out[23]:
[<matplotlib.lines.Line2D at 0x7f8b9d44f5c0>]
In [24]:
Ur[0, :] = Ut(0, x).real
Ui[0, :] = Ut(-0.5*dt, x).imag
#Ui[0, :] = U0.imag
schr_solve(Ur, Ui, dx, dt, nx, len(t))
In [9]:
#pcolormesh(sqrt(Ui**2+Ur**2)[:100,:], rasterized=True)
pcolormesh(Ui[:200,:], rasterized=True)
Out[9]:
<matplotlib.collections.QuadMesh at 0x7f8b9d070fd0>
In [10]:
# http://nbviewer.ipython.org/github/empet/Math/blob/master/DomainColoring.ipynb

def Hcomplex(z):# computes the hue corresponding to the complex number z
    H=np.angle(z)/(2*np.pi)+1
    return np.mod(H,1)
def g(x):
    return (1.0-1/(1+x**2))**0.2

U = Ur[:200,:] + 1j*Ui[:200,:]

from matplotlib.colors import hsv_to_rgb
H = Hcomplex(U)
V=g(abs(U)*0.5)
S = ones_like(H, float)

HSV = dstack((H,S,V))# V**0.2>V for V in[0,1];this choice  avoids too dark colors
RGB=hsv_to_rgb(HSV)
imshow(RGB)
Out[10]:
<matplotlib.image.AxesImage at 0x7f8b9cccef98>
In [11]:
from matplotlib import animation
from JSAnimation.IPython_display import display_animation
In [12]:
fig = figure()
ax = axes(xlim=(0,maxx),ylim=(-3,3))
line1, = ax.plot([],[],lw=1)
line2, = ax.plot([],[],lw=1)
line3, = ax.plot([],[],lw=2)
line4, = ax.plot([],[],lw=2)

def animate(data):
    line1.set_data(x,data[:len(x)])
    line2.set_data(x,data[len(x):])
    line3.set_data(x,sqrt(data[:len(x)]**2 + data[len(x):]**2))
    return line1,
In [13]:
Ui_int = 0.5*(Ui[:-1,:]+Ui[1:,:])
anim = animation.FuncAnimation(fig, animate, frames=hstack((Ui_int, Ur[:-1,:]))[:200:2,:], interval=100)
display_animation(anim, default_mode='loop')
Out[13]:


Once Loop Reflect