! Time independent quantum mechanics - square well with shooting method
! Program to accompany "Computational Physics" by N. Giordano
! Copyright Prentice Hall 1997
program square_well
   option nolet
   dim psi(-1 to 2000)   ! wave function (only need to worry about real part)
   call initialize(psi,dx,V,e_guess,de)
   call calculate(psi,dx,V,e_guess,de)
end
! initialize variables
! psi = wave function  dx = spatial step   square well runs from -xmax to xmax
! V = potential outside well   parity = +1 (even) or -1 (odd)
! e = energy and de = energy step 
sub initialize(psi(),dx,V,e,de)
   dx = 0.01
   V = 1e5
   e = 3
   de = -.6
   set window -1.5,1.5,-1.5,1.5    ! set up for plotting
   clear
   plot -1.5,0;1.5,0
   plot -1,-1;-1,1
   plot 1,-1;1,1
end sub
! psi() = wave function, dx = spatial grid size, Vmax is potential outside box
! energy is current best guess for E, de is amount E is changed in hunting
! for a solution
sub calculate(psi(),dx,Vmax,energy,de)
   declare def potential     ! potential energy function
   psi(0) = 1                ! search for an even parity solution
   psi(-1) = 1
   last_diverge = 0          ! use to keep track of direction of last divergence
                             ! loop as we zero in on the proper value of E
   print "Energy (hit any key to go to the next trial energy)"
   do                        
      print energy
      for i = 0 to size(psi)-1    ! integrate from x=0 to 1 
         psi(i+1) = 2 * psi(i) - psi(i-1) - 2 * (energy - potential(i*dx,Vmax))*dx^2*psi(i)
         if abs(psi(i+1)) > 2 then exit for   ! psi is diverging so stop now
      next i
      call display(psi,i,dx)           ! display this estimate for psi
      get key z                   ! wait for a key stroke before continuing
      if chr$(z) = "q" then exit do ! exit if the current solution is close enough
      if psi(i+1) > 0 then
         diverge = +1
      else
         diverge = -1
      end if
      if diverge * last_diverge < 0 then de = -de / 2
      energy = energy + de
      last_diverge = diverge
   loop
end sub
def potential(x,V)   ! the potential is 0 inside the box and V outside 
  if abs(x) <= 1 then   ! walls are at x = +1 and -1
      potential = 0
  else
      potential = V
  end if
end def
! display the wave function
sub display(psi(),max,dx)
   for i = -max to 0       ! assume an even parity solution
      plot i*dx,psi(-i);
   next i
   for i = 1 to max-1
      plot i*dx,psi(i);
   next i
   plot max*dx,psi(i)
end sub