%matplotlib inline
from pylab import *
from numpy import *

#physical constants
hbar = 1.0
m = 1.0

#discretization
nx = 100

minx = -1.
maxx = 1.

x = linspace(minx, maxx, nx)

#assemble the matrix Hamiltonian
from d2matrix import d2matrix
from scipy.sparse import dia_matrix

def make_H(V):
    nx = len(V)
    kinetic = -hbar**2/(2*m)*d2matrix(nx)
    potential = diag(V)
    H = kinetic + potential
    return H


V = zeros(nx)

H = make_H(V)

# calculate the eigenvectors
eigvals, eigvecs = linalg.eig(H)
eigvecs = array(eigvecs)


for i in arange(nx):
    plot(x, eigvecs[:,i]+eigvals[i]*150)
ylim([0, 10])

def solve(H):
    eigvals, eigvecs = linalg.eig(H)
    
    I = eigvals.argsort()
    eigvals = eigvals[I]
    eigvecs = array(eigvecs)[:,I]
    return eigvals, eigvecs

V = x**2
H = make_H(V)
eigvals, eigvecs = solve(H)

figure(figsize(5,6))
for i in arange(10):
    plot(x, eigvecs[:,i]*.05+eigvals[i])
plot(x, V, lw=3)
ylim([0, 0.3])

#discretization
nx = 100

minx = -1.
maxx = 1.

x = linspace(minx, maxx, nx)

V = zeros(nx)
a1 = -0.5
a2 = 0.5
asize = 0.2
apot = -0.07
V[(x>a1-asize) & (x<a1+asize) | (x>a2-asize) & (x<a2+asize) ] = apot

eigvals, eigvecs = solve(make_H(V))

figure(figsize(6,6))
for i in arange(6):
    plot(x, eigvecs[:,i]*.05+eigvals[i])
plot(x, V, lw=2)
set_printoptions(precision=16)
print(eigvals[:6])

#ylim([-1, -0.6])

coeffs = zeros(nx)
coeffs[0] = 1/sqrt(2.)
coeffs[1] = 1/sqrt(2.)
E_mean = dot(coeffs*eigvals, coeffs)

def psi_t(eigvecs, eigvals, coeffs, t):
    phase = exp(eigvals*t/1j*hbar)
    #print phase
    
    psi = zeros(len(coeffs), dtype=complex)
    for i in arange(len(coeffs)):
        if coeffs[i] != 0.:
            psi += eigvecs[:,i]*phase[i]
    return psi

from matplotlib import animation
from JSAnimation.IPython_display import display_animation

fig = figure()
ax = axes(xlim=(minx,maxx))
line1, = ax.plot([],[],lw=1)
ax.plot(x, V)


def animate(data):
    line1.set_data(x, abs(psi_t(eigvecs, eigvals, coeffs, 2000000.*data))*.1+E_mean)

    return line1,

anim = animation.FuncAnimation(fig, animate, frames=arange(20), interval=100)
display_animation(anim, default_mode='loop')