%matplotlib inline
from pylab import *
from numpy import *

#physical constants
hbar = 1.0
m = 1.0

#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

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]

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)

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



#plot(U0.real)
#plot(U0.imag)
plot(Ut(0, x).imag)
plot(Ut(0, x).real)
plot(abs(Ut(0, x)))

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))

#pcolormesh(sqrt(Ui**2+Ur**2)[:100,:], rasterized=True)
pcolormesh(Ui[:200,:], rasterized=True)

# 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)

from matplotlib import animation
from JSAnimation.IPython_display import display_animation

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,

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')

maxx = 5.
maxt = 0.08
nx = 500

sigma = 0.15
k0 = 40.
x0 = sigma*4.
E = (hbar**2/2.0/m)*(k0**2+0.5/sigma**2)

Ur, Ui, dx, dt, x, t = schr_init(maxx, maxt, E*1.5, nx, 0.7)

V = zeros_like(x)
xb = maxx*0.5
sb = 0.02
V = exp(-(x-xb)**2/(2*sb))*E*2.2
#V[(x>maxx*0.5) & (x<maxx*0.6)] = E*1.5

print(E)
print(len(t))
print(len(x))
Ur[0, :] = Ut(0, x).real
Ui[0, :] = Ut(-0.5*dt, x).imag
schr_solve(Ur, Ui, dx, dt, nx, len(t))

Ui_int = 0.5*(Ui[:-1,:]+Ui[1:,:])
#ax.cla()
ax.set_xlim([0,maxx])
line4.set_data(x, (V-E)/V.max())
anim = animation.FuncAnimation(fig, animate, frames=hstack((Ui_int, Ur[:-1,:]))[::20,:], interval=100)
display_animation(anim, default_mode='loop')

anim.save("animation.gif", writer='imagemagick', fps=10)