#! /usr/bin/env python
# Author: Are Raklev
# Date: 01.02.13
# Email: ahye@fys.uio.no
# Note the lack of units in wavefunctions. This is just a test of principle.

#  Numerical and plotting libraries (+ time)
from pylab import *
import time

# Constants that can be changed
k0 = 2.0           # Central wavenumber of packet
dk = 0.05          # Wavenumber difference between waves in superposition
n = 3             # Gives 2n+1 waves in wavepacket

# Define real and imaginary part of wavefunction
def psikw_r(x,t,k,w):
    return cos(k*x-w*t)
def psikw_i(x,t,k,w):
    return sin(k*x-w*t)
# Define sum of wavefunctions
def psi_sum_r(x,t):
    psi = 0.0
    for i in arange(-n,n):
        k = k0+dk*i
        w = 0.5*k**2
        psi += psikw_r(x,t,k,w)
    return psi
def psi_sum_i(x,t):
    psi = 0.0
    for i in arange(-n,n):
        k = k0+dk*i
        w = 0.5*k**2
        psi += psikw_i(x,t,k,w)
    return psi

# Set pylab in interactive mode
ion()

# Constants
N = 1000     #  Number of spatial points.
dx   = 0.1e0    #  Spatial resolution
x    = dx*linspace(-N,N,N)        #  Spatial axis.

# Draw initial state
lineP, = plot(x,psi_sum_r(x,0)**2+psi_sum_i(x,0)**2,'k',label='Probability')
lineR, = plot(x,psi_sum_r(x,0),'b',alpha=0.7,label='Real')
lineI, = plot(x,psi_sum_i(x,0),'g',alpha=0.7,label='Imaginary')
legend(loc='lower right')
draw()

tstart = time.time()               # for time profiling
for t in arange(1,200):
    lineR.set_ydata(psi_sum_r(x,t))  # update the data
    lineI.set_ydata(psi_sum_i(x,t))  # update the data
    lineP.set_ydata(psi_sum_r(x,t)**2+psi_sum_i(x,t)**2)  # update the data
    draw()               # redraw the canvas 
print 'FPS:' , 200/(time.time()-tstart)

# So the windows don't auto-close at the end if run outside ipython
ioff()
show()