# +----------------------------------------------------+
# |                  Lekce 3 - demo 2                  |
# +----------------------------------------------------+
#  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 <a,b> x <0,T> 
#  note: units are such that m=0.5, and planck h=1.0
#  THIS VERSION: test for nonzero V - barrier

import numpy as np
import scipy.linalg as linalg
import matplotlib.pyplot as plt

#interval x \in <a,b>
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 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%10==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/20.0,'b')
      plt.show()
    StepCN(u)