#!/usr/bin/env python3

"""
Animation of the Schrodinger equation for a particle in a 1D box

    -∂²/∂x² ψ(x,t) = i ∂/∂t ψ(x,t)          (natural units)

for a non-eigenstate initial state ψ, with boundary conditions

    ψ(0,t) = ψ(1,t) = 0.

The solution to the above differential equation is found by separation of
variables, and is as follows:

    ψ(x,t) = ∑(C_n exp(-i E_n t) sin(√(2 E_n) x)

Where C_n is arbitrary and the above represents a superposition of eigenstates,
and E_n is the n-th energy level:

    E_n = n²π²/2

"""

import scipy as sp
import scipy.linalg as la
import matplotlib.pyplot as plt
import matplotlib.animation as anim


def energy_level(n):
    return (n**2 * sp.pi**2) / 2


def wavefunction(x, t, coeff):
    n = sp.arange(len(coeff))
    E_n = energy_level(n)
    return (
        sp.sum(coeff * sp.exp(1j * t * E_n) * sp.sin(sp.sqrt(2*E_n) * x))/2
        )


coefficients = sp.rand(5)
coefficients = coefficients / la.norm(coefficients)

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(0, 1))
line, = ax.plot([], [])


def init():
    line.set_data([], [])
    return line,


def animate(t):
    xs = sp.linspace(0.0, 1.0, 200)
    ys = sp.array([wavefunction(x, 0.003 * t, coefficients) for x in xs])
    line.set_data(xs, ys * sp.conj(ys))
    return line,

anim = anim.FuncAnimation(fig, animate, init_func=init, frames=1000,
                          interval=16, blit=False)

plt.show()