# Script to simulate tunnelling in 1-D
# Isaac Lenton
# August 2017
import numpy as np
from numpy import fft
import matplotlib.pyplot as plt
import matplotlib.animation as animation

barrierWidth = 0.12;
barrierHeight = 2.1;
animLength = 1000;

xmax = 10;
dt = 1.0e-2;                            # Time step size
N = 512;                                # Simulation spatial resolution
x = np.arange(-N/2, N/2)*2.0*xmax/N;      # Position axis
p = np.arange(-N/2, N/2)*2.0*np.pi/xmax;  # Momentum axis

# Calculate potential
potential = 0.25*x**2 + barrierHeight*np.exp(-(x/barrierWidth)**2.0);

# Calculate initial wave-function
wavefunction = np.sqrt(np.sqrt(1/(2.0*np.pi)))*np.exp(-0.25*(x+4)**2.0);

# Calculate the propagators
kprop = np.exp(-1j*dt*p**2);
xprop = np.exp(-1j*dt*potential);

# Setup plot
fig, ax1 = plt.subplots()
line1, = ax1.plot(x, potential, 'b')
ax1.set_ylim(0.0, 25.0)
ax1.set_xlim(-xmax, xmax)
ax1.set_xlabel("Position")

ax2 = ax1.twinx()
line2, = ax2.plot(x, np.abs(wavefunction)**2.0, 'r')
ax2.set_ylim(0.0, 1.0)

fig.legend((line1, line2), ('Potential', '|Wavefunction|^2'))

def animate(step_number):

    global wavefunction

    # Evolve the wave-function
    kphi = fft.fftshift(fft.fft(fft.fftshift(wavefunction)))*kprop;
    wavefunction = fft.fftshift(fft.ifft(fft.fftshift(kphi)))*xprop;

    # Update the plot
    line2.set_data(x, np.abs(wavefunction)**2.0)
    return line2,

ani = animation.FuncAnimation(fig, animate, xrange(1, animLength),
    interval=25, blit=True, repeat=False)
#ani.save('output.mp4')
plt.show()