In [4]:
from pylab import *
import pandas as pd
from scipy.special import hermite
from scipy.misc import factorial


%matplotlib inline
# the harmonic oscillator potential with hbar = m = omega =1

def V(x):
    return x ** 2
# computational grid running from -5 to 5 with 512 grid points
n=4   # mean energy [( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 )]/2
E_mean=(( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 ))/3
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)   # size of spartial grid spacing

V0 = 1.

dt = 0.25 * dx ** 2
# temporal step size
# create Hamiltonian matrix and calculate its eigenvectors and
# eigenvalues
# use finite differences to approximate derivatives to 2 nd order

def Crank_Nicolson(psi, x, t, dt, V):
    print(np.absolute(sum(psi * dt)))
    H =(-0.5 * (diag(ones(N - 1), -1) - diag(2 * ones( N )) +
        diag(ones(N - 1), + 1)) / (dx ** 2) + diag(V(x))) * dt
    return solve((np.eye(N) + 1j / 2 * H),(np.eye(N) - 1j / 2 * H).dot(psi))

def Psi(n, x):
    H=hermite(n)
    return 1/sqrt(2**n*factorial(n))*1/pi**0.25*exp(-x**2/2)*H(x)

In [ ]:
psi = Psi(0, x) #+ Psi(1, x) + Psi(2, x)
for t in  np.arange(200) * dt:
    psi = Crank_Nicolson(psi, x, t, dt, V)
    plot(x ,psi.real , color='#266bbd', label=r'$|\Psi(x)|^2$')


0.00089853609218

In [236]:
psi0 = Psi(0, x) + Psi(1, x)
for t in  np.arange(0, 80, 20):
    psi = Psi(0, x) * np.exp(-1j * t / 2) + Psi(1, x) * np.exp(-1j * t *(1/ 2 + 1))
    plot(x, psi.real ** 2 + psi.imag ** 2 , color='#266bbd', label=r'$|\Psi(x)|^2$')



In [129]:
psi


Out[129]:
array([  2.03857416e-20,   2.98915701e-20,   4.37614847e-20,
         6.39670566e-20,   9.33558580e-20,   1.36034028e-19,
         1.97912947e-19,   2.87489112e-19,   4.16954849e-19,
         6.03777563e-19,   8.72941434e-19,   1.26012411e-18,
         1.81619086e-18,   2.61354198e-18,   3.75506259e-18,
         5.38672069e-18,   7.71527096e-18,   1.10330917e-17,
         1.57529706e-17,   2.24567419e-17,   3.19631781e-17,
         4.54225950e-17,   6.44484644e-17,   9.13001957e-17,
         1.29136557e-16,   1.82366438e-16,   2.57133521e-16,
         3.61984854e-16,   5.08791533e-16,   7.14014400e-16,
         1.00044111e-15,   1.39956605e-15,   1.95484542e-15,
         2.72614173e-15,   3.79578204e-15,   5.27680160e-15,
         7.32414204e-15,   1.01498372e-14,   1.40435707e-14,
         1.94004583e-14,   2.67585291e-14,   3.68492034e-14,
         5.06651589e-14,   6.95514145e-14,   9.53273654e-14,
         1.30450002e-13,   1.78231831e-13,   2.43131335e-13,
         3.31139423e-13,   4.50292720e-13,   6.11354119e-13,
         8.28713534e-13,   1.12157800e-12,   1.51554130e-12,
         2.04465058e-12,   2.75412236e-12,   3.70390589e-12,
         4.97334986e-12,   6.66730316e-12,   8.92407645e-12,
         1.19258134e-11,   1.59119771e-11,   2.11968554e-11,
         2.81922437e-11,   3.74367809e-11,   4.96338256e-11,
         6.57002671e-11,   8.68293171e-11,   1.14571139e-10,
         1.50936192e-10,   1.98527460e-10,   2.60709307e-10,
         3.41822733e-10,   4.47459276e-10,   5.84808872e-10,
         7.63100869e-10,   9.94162074e-10,   1.29312160e-09,
         1.67929947e-09,   2.17732479e-09,   2.81854019e-09,
         3.64276248e-09,   4.70048577e-09,   6.05563293e-09,
         7.78898551e-09,   1.00024511e-08,   1.28243629e-08,
         1.64160475e-08,   2.09799502e-08,   2.67696671e-08,
         3.41023061e-08,   4.33736909e-08,   5.50770222e-08,
         6.98257385e-08,   8.83814674e-08,   1.11688130e-07,
         1.40913477e-07,   1.77499558e-07,   2.23223951e-07,
         2.80273864e-07,   3.51335650e-07,   4.39702719e-07,
         5.49405356e-07,   6.85366571e-07,   8.53588803e-07,
         1.06137713e-06,   1.31760556e-06,   1.63303400e-06,
         2.02068481e-06,   2.49628910e-06,   3.07881463e-06,
         3.79108872e-06,   4.66053179e-06,   5.72001916e-06,
         7.00889123e-06,   8.57413500e-06,   1.04717627e-05,
         1.27684165e-05,   1.55432329e-05,   1.88900017e-05,
         2.29196626e-05,   2.77631832e-05,   3.35748688e-05,
         4.05361607e-05,   4.88599818e-05,   5.87956972e-05,
         7.06347607e-05,   8.47171260e-05,   1.01438506e-04,
         1.21258567e-04,   1.44710161e-04,   1.72409679e-04,
         2.05068652e-04,   2.43506686e-04,   2.88665861e-04,
         3.41626688e-04,   4.03625752e-04,   4.76075142e-04,
         5.60583773e-04,   6.58980711e-04,   7.73340574e-04,
         9.06011101e-04,   1.05964294e-03,   1.23722168e-03,
         1.44210220e-03,   1.67804521e-03,   1.94925604e-03,
         2.26042553e-03,   2.61677284e-03,   3.02409004e-03,
         3.48878825e-03,   4.01794481e-03,   4.61935138e-03,
         5.30156220e-03,   6.07394205e-03,   6.94671335e-03,
         7.93100138e-03,   9.03887703e-03,   1.02833959e-02,
         1.16786328e-02,   1.32397103e-02,   1.49828203e-02,
         1.69252371e-02,   1.90853203e-02,   2.14825063e-02,
         2.41372863e-02,   2.70711700e-02,   3.03066322e-02,
         3.38670414e-02,   3.77765677e-02,   4.20600699e-02,
         4.67429586e-02,   5.18510349e-02,   5.74103025e-02,
         6.34467529e-02,   6.99861226e-02,   7.70536213e-02,
         8.46736311e-02,   9.28693771e-02,   1.01662570e-01,
         1.11073020e-01,   1.21118228e-01,   1.31812952e-01,
         1.43168748e-01,   1.55193506e-01,   1.67890959e-01,
         1.81260189e-01,   1.95295137e-01,   2.09984103e-01,
         2.25309264e-01,   2.41246207e-01,   2.57763487e-01,
         2.74822218e-01,   2.92375703e-01,   3.10369117e-01,
         3.28739252e-01,   3.47414327e-01,   3.66313873e-01,
         3.85348718e-01,   4.04421050e-01,   4.23424595e-01,
         4.42244905e-01,   4.60759752e-01,   4.78839657e-01,
         4.96348533e-01,   5.13144462e-01,   5.29080596e-01,
         5.44006187e-01,   5.57767737e-01,   5.70210268e-01,
         5.81178703e-01,   5.90519347e-01,   5.98081459e-01,
         6.03718897e-01,   6.07291825e-01,   6.08668465e-01,
         6.07726865e-01,   6.04356670e-01,   5.98460881e-01,
         5.89957554e-01,   5.78781445e-01,   5.64885545e-01,
         5.48242512e-01,   5.28845946e-01,   5.06711503e-01,
         4.81877815e-01,   4.54407201e-01,   4.24386141e-01,
         3.91925506e-01,   3.57160521e-01,   3.20250462e-01,
         2.81378058e-01,   2.40748618e-01,   1.98588871e-01,
         1.55145514e-01,   1.10683501e-01,   6.54840542e-02,
         1.98424462e-02,  -2.59344578e-02,  -7.15308485e-02,
        -1.16624798e-01,  -1.60891228e-01,  -2.04004979e-01,
        -2.45643966e-01,  -2.85492347e-01,  -3.23243707e-01,
        -3.58604178e-01,  -3.91295487e-01,  -4.21057868e-01,
        -4.47652810e-01,  -4.70865614e-01,  -4.90507699e-01,
        -5.06418650e-01,  -5.18467970e-01,  -5.26556508e-01,
        -5.30617548e-01,  -5.30617548e-01,  -5.26556508e-01,
        -5.18467970e-01,  -5.06418650e-01,  -4.90507699e-01,
        -4.70865614e-01,  -4.47652810e-01,  -4.21057868e-01,
        -3.91295487e-01,  -3.58604178e-01,  -3.23243707e-01,
        -2.85492347e-01,  -2.45643966e-01,  -2.04004979e-01,
        -1.60891228e-01,  -1.16624798e-01,  -7.15308485e-02,
        -2.59344578e-02,   1.98424462e-02,   6.54840542e-02,
         1.10683501e-01,   1.55145514e-01,   1.98588871e-01,
         2.40748618e-01,   2.81378058e-01,   3.20250462e-01,
         3.57160521e-01,   3.91925506e-01,   4.24386141e-01,
         4.54407201e-01,   4.81877815e-01,   5.06711503e-01,
         5.28845946e-01,   5.48242512e-01,   5.64885545e-01,
         5.78781445e-01,   5.89957554e-01,   5.98460881e-01,
         6.04356670e-01,   6.07726865e-01,   6.08668465e-01,
         6.07291825e-01,   6.03718897e-01,   5.98081459e-01,
         5.90519347e-01,   5.81178703e-01,   5.70210268e-01,
         5.57767737e-01,   5.44006187e-01,   5.29080596e-01,
         5.13144462e-01,   4.96348533e-01,   4.78839657e-01,
         4.60759752e-01,   4.42244905e-01,   4.23424595e-01,
         4.04421050e-01,   3.85348718e-01,   3.66313873e-01,
         3.47414327e-01,   3.28739252e-01,   3.10369117e-01,
         2.92375703e-01,   2.74822218e-01,   2.57763487e-01,
         2.41246207e-01,   2.25309264e-01,   2.09984103e-01,
         1.95295137e-01,   1.81260189e-01,   1.67890959e-01,
         1.55193506e-01,   1.43168748e-01,   1.31812952e-01,
         1.21118228e-01,   1.11073020e-01,   1.01662570e-01,
         9.28693771e-02,   8.46736311e-02,   7.70536213e-02,
         6.99861226e-02,   6.34467529e-02,   5.74103025e-02,
         5.18510349e-02,   4.67429586e-02,   4.20600699e-02,
         3.77765677e-02,   3.38670414e-02,   3.03066322e-02,
         2.70711700e-02,   2.41372863e-02,   2.14825063e-02,
         1.90853203e-02,   1.69252371e-02,   1.49828203e-02,
         1.32397103e-02,   1.16786328e-02,   1.02833959e-02,
         9.03887703e-03,   7.93100138e-03,   6.94671335e-03,
         6.07394205e-03,   5.30156220e-03,   4.61935138e-03,
         4.01794481e-03,   3.48878825e-03,   3.02409004e-03,
         2.61677284e-03,   2.26042553e-03,   1.94925604e-03,
         1.67804521e-03,   1.44210220e-03,   1.23722168e-03,
         1.05964294e-03,   9.06011101e-04,   7.73340574e-04,
         6.58980711e-04,   5.60583773e-04,   4.76075142e-04,
         4.03625752e-04,   3.41626688e-04,   2.88665861e-04,
         2.43506686e-04,   2.05068652e-04,   1.72409679e-04,
         1.44710161e-04,   1.21258567e-04,   1.01438506e-04,
         8.47171260e-05,   7.06347607e-05,   5.87956972e-05,
         4.88599818e-05,   4.05361607e-05,   3.35748688e-05,
         2.77631832e-05,   2.29196626e-05,   1.88900017e-05,
         1.55432329e-05,   1.27684165e-05,   1.04717627e-05,
         8.57413500e-06,   7.00889123e-06,   5.72001916e-06,
         4.66053179e-06,   3.79108872e-06,   3.07881463e-06,
         2.49628910e-06,   2.02068481e-06,   1.63303400e-06,
         1.31760556e-06,   1.06137713e-06,   8.53588803e-07,
         6.85366571e-07,   5.49405356e-07,   4.39702719e-07,
         3.51335650e-07,   2.80273864e-07,   2.23223951e-07,
         1.77499558e-07,   1.40913477e-07,   1.11688130e-07,
         8.83814674e-08,   6.98257385e-08,   5.50770222e-08,
         4.33736909e-08,   3.41023061e-08,   2.67696671e-08,
         2.09799502e-08,   1.64160475e-08,   1.28243629e-08,
         1.00024511e-08,   7.78898551e-09,   6.05563293e-09,
         4.70048577e-09,   3.64276248e-09,   2.81854019e-09,
         2.17732479e-09,   1.67929947e-09,   1.29312160e-09,
         9.94162074e-10,   7.63100869e-10,   5.84808872e-10,
         4.47459276e-10,   3.41822733e-10,   2.60709307e-10,
         1.98527460e-10,   1.50936192e-10,   1.14571139e-10,
         8.68293171e-11,   6.57002671e-11,   4.96338256e-11,
         3.74367809e-11,   2.81922437e-11,   2.11968554e-11,
         1.59119771e-11,   1.19258134e-11,   8.92407645e-12,
         6.66730316e-12,   4.97334986e-12,   3.70390589e-12,
         2.75412236e-12,   2.04465058e-12,   1.51554130e-12,
         1.12157800e-12,   8.28713534e-13,   6.11354119e-13,
         4.50292720e-13,   3.31139423e-13,   2.43131335e-13,
         1.78231831e-13,   1.30450002e-13,   9.53273654e-14,
         6.95514145e-14,   5.06651589e-14,   3.68492034e-14,
         2.67585291e-14,   1.94004583e-14,   1.40435707e-14,
         1.01498372e-14,   7.32414204e-15,   5.27680160e-15,
         3.79578204e-15,   2.72614173e-15,   1.95484542e-15,
         1.39956605e-15,   1.00044111e-15,   7.14014400e-16,
         5.08791533e-16,   3.61984854e-16,   2.57133521e-16,
         1.82366438e-16,   1.29136557e-16,   9.13001957e-17,
         6.44484644e-17,   4.54225950e-17,   3.19631781e-17,
         2.24567419e-17,   1.57529706e-17,   1.10330917e-17,
         7.71527096e-18,   5.38672069e-18,   3.75506259e-18,
         2.61354198e-18,   1.81619086e-18,   1.26012411e-18,
         8.72941434e-19,   6.03777563e-19,   4.16954849e-19,
         2.87489112e-19,   1.97912947e-19,   1.36034028e-19,
         9.33558580e-20,   6.39670566e-20,   4.37614847e-20,
         2.98915701e-20,   2.03857416e-20])

In [181]:
import scipy.linalg as la

m = 3
# Create arrays and set values
ab = np.zeros((3,m))
b = 2*ones(m)
ab[0] = 9
ab[1] = 8
ab[2] = 7

# Fix end points
ab[0,1] = 2
ab[1,0] = 1
ab[1,-1] = 4
ab[2,-2] = 3
b[0] = 1
b[-1] = 3

return la.solve_banded ((1,1),ab,b)


  File "<ipython-input-181-6db26b14a74b>", line 19
    return la.solve_banded ((1,1),ab,b)
                                       ^
SyntaxError: 'return' outside function

In [171]:
H =(-0.5 * (diag(ones(N - 1), -1) - diag(2 * ones( N )) +
        diag(ones(N - 1), + 1)) / (dx ** 2) + diag(V(x))) * dt
'''
H = np.empty((3, N))
H[0, 1:] = -.5 * ones(N - 1)
H[2, :-1] = -.5 * ones(N - 1)
H[1, :] = 2 * ones( N )
H /= dx ** 2
H[1, :] += V(x)
H * dt
'''
diagonal(H, 1)


Out[171]:
array([-5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953, -5.10001953, -5.10001953, -5.10001953, -5.10001953,
       -5.10001953])

In [172]:
pd.DataFrame(H)


Out[172]:
0 1 2 3 4 5 6 7 8 9 ... 502 503 504 505 506 507 508 509 510 511
0 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
1 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
2 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
3 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
4 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
5 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
6 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
7 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
8 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
9 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
10 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
11 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
12 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
13 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
14 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
15 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
16 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
17 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
18 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
19 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
20 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
21 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
22 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
23 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
24 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
25 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
26 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
27 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
28 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
29 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
482 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
483 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
484 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
485 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
486 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
487 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
488 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
489 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
490 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
491 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
492 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
493 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
494 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
495 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
496 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
497 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
498 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
499 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
500 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
501 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
502 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
503 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
504 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000
505 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000 0.000000
506 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000 0.000000
507 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000 0.000000
508 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000 0.000000
509 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020 0.000000
510 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039 -5.100020
511 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 ... 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 0.000000 -5.100020 10.200039

512 rows × 512 columns


In [1]:
from pylab import *
from scipy.linalg import solve_banded
from scipy.special import hermite
from scipy.misc import factorial

close('all')

# the harmonic oscillator potential with hbar=m=omega=1
def V(x):
    return 0.5*x**2

# eigen functions of the quantum harmonic oscillator
# use Python's built-in hermite function
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)

# get the default colors
prop_cycle=plt.rcParams['axes.prop_cycle']
colors=prop_cycle.by_key()['color']

n=4   # mean energy [( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 )]/2
E_mean=(( (n+1) +1/2 ) + (n+1/2) + ( (n+1) +1/2 ))/3
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)   # size of spartial grid spacing

dt=2**linspace(-12, 0, 13)
error=empty_like(dt)
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_
    Psi_exat=(Psi_h(x, n-1, t)+Psi_h(x, n, t)+Psi_h(x, n+1, t))/sqrt(3)
    error[i]=norm(Psi-Psi_exat)*dx
loglog(dt, error, 'o')
# global error is O(dt^2) because single-step error is O(dt^3) and
# smaller step size requires more steps
loglog(dt, dt**2, '-', label=r'$\sim\Delta t^2$')
xlabel(r'$\Delta t$')
ylabel('error')
legend()
tight_layout()
show()



In [3]:
p_max, x1, x0


Out[3]:
(3.2145502536643185, 10.0, -10.0)

In [ ]: