from scipy import *
from scipy import integrate, optimize
from pylab import *

def Schrod_deriv(y, r, l, E, Z):
    "Given y=[u,u'] returns dy/dr=[u',u'']"
    du2 = (l*(l+1)/r**2-2*Z/r-E)*y[0]
    return [y[1],du2]

def Shoot(E,R,l,Z):
    y0=[0.0,-1e-7]
    Rb=R[::-1]
    y = integrate.odeint(Schrod_deriv, y0, Rb, args=(l,E ,Z))[:,0]
    norm = integrate.simps(Rb, y**2)
    f0 = y[-1]/sqrt(norm)
    f1 = y[-2]/sqrt(norm)
    final = f0 + (f1-f0)*(0.0-Rb[-1])/(Rb[-2]-Rb[-1])
    #print 'f0=', f0, 'f1=', f1, 'extrapolated=', final
    return final
    


(l, E, Z) = (0, -1., 1.)
MaxSteps=1000
nmax=5

R = logspace(-7,2.2,500)
    
E0 = -1.2*Z**2
dE = 0.1*Z**2
Eb=[]  # bound states
u0 = Shoot(E0,R,l,Z)
for i in range(MaxSteps):
    E0 += dE
    u1 = Shoot(E0,R,l,Z)
    print E0, u1, dE
    if u0*u1<0.0:
        Ebound = optimize.brentq(Shoot, E0-dE, E0, args=(R,l,Z))
        Eb.append( Ebound )
        print 'Found bound state at E=', Ebound
        if len(Eb)>=nmax or Ebound>0: break
    u0=u1
    dE = dE/1.091
print Eb
