import scipy as sp
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# animation of the probability density of the wavefunction over the course
# of time
def probabilityDensity(x, t, V, psi):
    
    # convert to the probability density
    Nt = len(t)
    rho = sp.real(sp.conjugate(psi)*psi)

    # set the first frame properties and grab the line handles
    fig, ax = plt.subplots()

    line1, line2, line3, line4 = ax.plot(x, rho[:,1], 'k',
            x, sp.real(psi[:,1]), 'b:',
            x, sp.imag(psi[:,1]), 'r:',
            x, V, 'm--',
            linewidth=2.0)

    ax.set_xlabel("Position")
    ax.set_ylabel("Probability Density")
    ax.set_ylim([-rho.max(), rho.max()])
    ax.set_xlim([min(x), max(x)])

    # the animation function, to be called repeatedly
    def animate(i):
        # set the new data each frame
        line1.set_ydata(rho[:,i])
        line2.set_ydata(sp.real(psi[:,i])) 
        line3.set_ydata(sp.imag(psi[:,i]))
        return line1, line2, line3

    # the initialization function, useful when blit=True
    def init():
        line1.set_ydata(sp.ma.array(x, mask=True))
        line2.set_ydata(sp.ma.array(x, mask=True))
        line3.set_ydata(sp.ma.array(x, mask=True))
        return line1, line2, line3

    # perform the animation
    ani = animation.FuncAnimation(fig, animate, sp.arange(1,Nt), 
            init_func=init, interval=25, blit=True)
    plt.show()
