# +----------------------------------------------------+
# | Lekce 3 - demo 1 |
# +----------------------------------------------------+
# FINITE DIFFERENCE METHODS for initial value problem
#
# PROBLEM: to solve initial value problem for Schrodinger equation
# dU/dt = i d^2U/dx^2 + iV(x)U on (x,t) in x <0,T>
# note: units are such that m=0.5, and planck h=1.0
# THIS VERSION: test for V=0
import numpy as np
import scipy.linalg as linalg
import matplotlib.pyplot as plt
#interval x \in
a=-10.0
b= 10.0
#wave-packet parameters:
p0=5.0
v0=2*p0
x0=-5.0
delta_x0=1.0
delta_p0=0.5/delta_x0
#wave function squared for exact free particle solution
def E_wfc2(xx,tt):
s=np.sqrt(delta_x0**2+(2*delta_p0*tt)**2)
xt=x0+v0*tt
return np.exp(-2*((xx-xt)/2/s)**2)/np.sqrt(2*np.pi)/s
#Parameters of run:
Nx=4000
h=(b-a)/Nx
T=2.0
dt=0.005
sigm=dt/h**2
def V(x):
return 0.0 #10.0*np.exp(-2*x**2)
#Construction of grigpoints and initial condition:
x=np.linspace(a,b,Nx+1)
nrm=(np.pi*2)**0.25*np.sqrt(delta_x0)
u=np.exp(1j*p0*x[:]-((x[:]-x0)/(2*delta_x0))**2)/nrm
u[0]=0+0j
u[-1]=0+0j
v=np.array([V(x[i]) for i in range(Nx+1)])
#Construction of matrix to be inverted
Mband=np.zeros([3,Nx-1],dtype=complex)
for i in range(1,Nx):
Mband[1,i-1]=2+2j*sigm+1j*dt*v[i]
Mband[0,i-1]=-1j*sigm
Mband[2,i-1]=-1j*sigm
#One step of Crank-Nicholson method
def StepCN(u):
u[1:-1]=u[1:-1]*(2-2j*sigm-1j*dt*v[1:-1])+1j*sigm*(u[0:-2]+u[2:])
u[1:-1]=linalg.solve_banded((1,1),Mband,u[1:-1])
t=0.0
#for n in range(int(T/dt)):
for n in range(400):
t=n*dt
if (n%40==0):
plt.text(0.0, 0.5,'t='+str(t)[:7])
eu=np.array([np.sqrt(E_wfc2(x[i],t)) for i in range(Nx+1)])
plt.plot(x,eu,'b')
plt.plot(x,np.real(u),'g')
plt.plot(x,np.abs(u),'r')
plt.plot(x,v,'b')
plt.show()
StepCN(u)