#!/usr/bin/env python
from pylab import *

N = 500
xmax = 5

x = linspace(-xmax,xmax,N)
dx = x[1]-x[0]

T = -1/2.*(eye(N,k=-1) - 2*eye(N) + eye(N,k=1))/dx/dx
V = diag(x*x/2.)

plot(x, diag(V), "k-")

H = T + V
E, psi = linalg.eigh(H)

o = 0
twinx()
for p,e in zip(psi.T, E)[:10]:
    prob = abs(p)**2
    plot(x, o+prob, label="E={:.4}".format(e))
    o += 1.1*amax(prob)
legend()

savefig("osc.png")
show()