In [20]:

    

from pylab import *
from scipy.linalg import solve_banded
from scipy.special import hermite
from scipy.misc import factorial

%matplotlib inline

def V(x):
    return 1.5 * x ** 2

def Psi_h(x, n, t=0):
    H=hermite(n)
    return 1/sqrt(2**n*factorial(n))*1/pi**0.25*exp(-x**2/2)*H(x)*exp(-1j*(n+0.5)*t)

'''
N = 256
x0, x1 = -5., 5.
x = linspace(x0, x1, N)
dx = (x1 - x0) / (N-1)
dt = 1. / 64 * 1
'''
n=1   # mean energy [( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 )]/2
E_mean= 3 / 2
p_max=sqrt(2*E_mean)
dx=1/512*pi/p_max

# computational grid depends on n
x0, x1=-sqrt(n)*5, sqrt(n)*5
N=int(round((x1-x0)/dx))
x=linspace(x0, x1, N)
dx=(x1-x0)/(N-1)

Psi = Psi_h(x, 0, t=0)

dp = 2 * pi / (N * dx)
p = linspace(0, (N - 1) * dp, N)
p[p > .5 * N * dp] - N * dp

U_1 = exp(-.5 * 1j * V(x) * dt)
U_2 = exp(-1j * .5 * p ** 2 * dt)
    
H=zeros((3, N))
H[0, 1: ]=-0.5/dx**2
H[1, :  ]=+1.0/dx**2
H[2, :-1]=-0.5/dx**2
Dp=+0.5j*dt*H
Dp[1, :]+=1
Dm=-0.5j*dt*H 
Dm[1, :]+=1

for k  in range(512 // 1):
    Psi *= U_1
    #Psi = fft(Psi)
    #Psi *= U_2
    #Psi = ifft(Psi)

    b=Dm[1, :]*Psi
    b[1:]+=Dm[0, 1:]*Psi[:-1]
    b[:-1]+=Dm[2, :-1]*Psi[1:]
    Psi=solve_banded((1, 1), Dp, b)    
    Psi *= U_1
    plot(np.absolute(Psi))


    

In [5]:

    

for i, dt_ in enumerate(dt):
    # create Hamiltonian matrix, use finite differences to approximate
    # derivatives to 2nd order, store banded matrices as 3 x N matrices
    H=zeros((3, N))
    H[0, 1: ]=-0.5/dx**2
    H[1, :  ]=+1.0/dx**2 + V(x)
    H[2, :-1]=-0.5/dx**2
    Dp=+0.5j*dt_*H
    Dp[1, :]+=1
    Dm=-0.5j*dt_*H
    Dm[1, :]+=1

    # propagate wave function
    Psi=(Psi_h(x, n-1)+Psi_h(x, n)+Psi_h(x, n+1))/sqrt(3)
    t=0
    while t<4:
        b=Dm[1, :]*Psi
        b[1:]+=Dm[0, 1:]*Psi[:-1]
        b[:-1]+=Dm[2, :-1]*Psi[1:]
        Psi=solve_banded((1, 1), Dp, b)
        t+=dt_


    

    




---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-e3d1f572cc94> in <module>()
----> 1 for i, dt_ in enumerate(dt):
      2     # create Hamiltonian matrix, use finite differences to approximate
      3     # derivatives to 2nd order, store banded matrices as 3 x N matrices
      4     H=zeros((3, N))
      5     H[0, 1: ]=-0.5/dx**2

TypeError: 'float' object is not iterable

In [18]:

    

n=4   # mean energy [( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 )]/2
E_mean=1/2
p_max=sqrt(2*E_mean)
dx=1/512*pi/p_max


    

In [11]:

    

plot((Psi.real))


    

    
Out[11]:





[<matplotlib.lines.Line2D at 0x7fd03765ef60>]






    

In [15]:

    

dx


    

    
Out[15]:





0.0392156862745098

In [16]:

    

dx=1/512*pi/p_max


    

In [19]:

    

dx


    

    
Out[19]:





0.0061359231515425647

In [ ]: