Particle in a box using Schrodinger Equation

#PARTICLE  INSIDE A BOX using RK method
from pylab import*
def f1(x,y,z):
    return(z)
def f2(x,y,z):
    return(y)
v=0  #potential inside well
m=1  #mass
y0=0
z0=1
hbar=1
n=int(input('enter the quantum no(n):'))
a=float(input('enter size of box:'))
e=pi*pi*n*n*hbar*hbar/(2*m*a*a)  #Energy
fn=2*m*(v-e)/(hbar*hbar)
h=0.005
x0=0
x=[]
yp=[] 
y=[]
z=[]
while x0<=a:
    k1=h*f1(x0,y0,z0)
    l1=h*f2(x0,y0,z0)*fn
    k2=h*f1(x0+h/2,y0+k1/2,z0+l1/2)
    l2=h*f2(x0+h/2,y0+k1/2,z0+l1/2)*fn
    k3=h*f1(x0+h/2,y0+k2/2,z0+l2/2)
    l3=h*f2(x0+h/2,y0+k2/2,z0+l2/2)*fn
    k4=h*f1(x0+h,y0+k3,z0+l3)
    l4=h*f2(x0+h,y0+k3,z0+l3)*fn
    k=(k1+2*k2+2*k3+k4)/6
    l=(l1+2*l2+2*l3+l4)/6
    x1=x0+h
    y1=y0+k
    z1=z0+l
    x0=x1
    y0=y1
    z0=z1
    psi=y0*y0
    x.append(x0)
    y.append(y0)
    yp.append(psi*10)
    z.append(z0)

print'energy e=',e
title('particle in a box')
grid(True)
plot(x,y)
plot(x,yp)
legend(['wave function','probability density'])
show()




Output