from visual.graph import *
from math import *

pic=gdisplay()
wave=gcurve(gdisplay=pic,color=color.blue)

# number of pts
N=10000
# integrate finite distance in x
BOX=3.0

dx=BOX/N
# initial guess at energy
E=2.5

phi=zeros((N),Float)
p=zeros((N),Float)
phi[0]=1.0
p[0]=0.0
# setting units in which m/hbar^2 = 1

# harmonic oscillator
def V(x):
    return 0.5*x*x

while(1):
    
    E=input("Input a trial energy ")
    print "energy is ", E
    wave.visible=False
    del wave
    pic.visible=False
    del pic
    
    
    for i in range(1,N):
        x=(i-1)*dx
        phi[i]=phi[i-1]+p[i-1]*dx+0.5*dx*dx*2*(V(x)-E)*phi[i-1]
        p[i]=p[i-1]+0.5*dx*2*(V(x)-E)*(phi[i]+phi[i-1])

    # normalize
    sum=0.0
    for i in range(0,N):
        sum+=phi[i]*phi[i]*dx

    pic=gdisplay()
    wave=gcurve(gdisplay=pic,color=color.blue)
    
    # plot wavefunction
    for i in range(0,N):
        wave.plot(pos=(i*dx,phi[i]/sqrt(sum)))