! wave packet propagation in one dimension
! Program to accompany "Computational Physics" by N. Giordano
! Copyright Prentice Hall 1997
program propagate
   option nolet
   library "sgfunc*", "sglib*"
!  use two dimensional arrays to store real and imaginary parts of psi, e, and f
   dim psi_old(0 to 2000,2),psi_new(0 to 2000,2),e(0 to 2000,2),f(0 to 2000,2)
   t = 0
   call initialize(psi_old,max,dx,dt)
   call display(psi_old,max,dx)
   n_display = 30   ! display psi after every 30*dt 
   do
      for i = 1 to n_display
         call calculate(psi_old,psi_new,e,f,max,dx,dt,2*dx*dx/dt,dx*dx)
      next i
      call display(psi_old,max,dx)
      t = t + n_display * dt
      if key input then
         get key z
         c$ = chr$(z)
         if c$ = "c" then
            clear
            call axis
         end if
         if c$ = "t" then print t;
         if c$ = "n" then call check_normalization(psi_old,max)
      end if
   loop 
end
! initialize variables
sub initialize(psi(,),max,dx,dt)
   max = 2000 
   dx = 1/max
   dt = 2 * dx^2   ! lambda = 1
   call init_psi_packet(psi(,),max,dx)
   clear
   set window 0,1,-0.2,40
   call axis
end sub
sub axis
   plot 0,0;1,0
end sub
! construct a traveling gaussian wave packet
sub init_psi_packet(p(,),m,dx)
   option angle radians
   x_0 = 0.4   ! wave packet is centered here
   k = 500     ! this is the wave vector
   for i = 0 to m
      a = exp(-(i*dx - x_0)^2/.001)    ! a gaussian packet
      p(i,1) = a * cos(k*i*dx)         ! real part of psi
      p(i,2) = a * sin(k*i*dx)         ! imaginary part of psi
   next i
   call normalize(p,m,dx)
end sub
! normalize psi - part of initialization
sub normalize(p(,),m,dx)
   sum = 0
   for i = 0 to m
      sum = sum + p(i,1)*p(i,1) + p(i,2)*p(i,2)
   next i
   sum = sqr(sum*dx)
   for i = 0 to m
      p(i,1) = p(i,1) / sum
      p(i,2) = p(i,2) / sum
   next i
end sub
! first calculate the new e and f factors, the update psi
sub calculate(p_old(,),p_new(,),e(,),f(,),max,dx,dt,lambda,dx2)
   declare def potential
   call calc_ef(p_old,e,f,max,dx,dt,lambda,dx2)
   call calc_psi(p_old,p_new,e,f,max)
   for i = 0 to max
      p_old(i,1) = p_new(i,1)
      p_old(i,2) = p_new(i,2)
   next i
end sub
! calculate the new wave function
sub calc_psi(p_old(,),p_new(,),e(,),f(,),max)
   declare def potential
   emod = e(max-1,1)^2 + e(max-1,2)^2
   p_new(max-1,1) = -(f(max-1,1)*e(max-1,1) + f(max-1,2)*e(max-1,2)) / emod
   p_new(max-1,2) = -(f(max-1,2)*e(max-1,1) - f(max-1,1)*e(max-1,2)) / emod
   for i = max-1 to 1 step -1
      emod = e(i,1)^2 + e(i,2)^2
      p_new(i,1) = ((p_new(i+1,1) - f(i,1))*e(i,1) + (p_new(i+1,2) - f(i,2))*e(i,2)) / emod
      p_new(i,2) = ((p_new(i+1,2) - f(i,2))*e(i,1) - (p_new(i+1,1) - f(i,1))*e(i,2)) / emod
   next i
end sub
! calculate the new e and f factors
sub calc_ef(p(,),e(,),f(,),max,dx,dt,lambda,dx2)
   declare def potential
   e(1,1) = 2 + 2 * dx2 * potential(dx) 
   e(1,2) = - 2 * lambda
   f(1,1) = -p(2,1) + (2*dx2*potential(dx) + 2) * p(1,1) - 2 * lambda * p(1,2) - p(0,1)
   f(1,2) = -p(2,2) + (2*dx2*potential(dx) + 2) * p(1,2) + 2 * lambda * p(1,1) - p(0,2)
   for i = 2 to max-1
      emod = e(i-1,1)^2 + e(i-1,2)^2
      e(i,1) = 2 + 2*dx2*potential(i*dx) - e(i-1,1) / emod
      e(i,2) = - 2 * lambda + e(i-1,2) / emod
      f(i,1) = -p(i+1,1) + (2*dx2*potential(i*dx) + 2) * p(i,1) -  2 * lambda * p(i,2) - p(i-1,1) + (f(i-1,1)*e(i-1,1) + f(i-1,2)*e(i-1,2)) / emod
      f(i,2) = -p(i+1,2) + (2*dx2*potential(i*dx) + 2) * p(i,2) +  2 * lambda * p(i,1) - p(i-1,2) + (f(i-1,2)*e(i-1,1) - f(i-1,1)*e(i-1,2)) / emod
   next i
end sub
sub display(p(,),m,dx)   ! display psi*psi
   for i = 0 to m-1
      plot i*dx,p(i,1)^2+p(i,2)^2;
   next i
   plot m*dx,p(m,1)^2+p(m,2)^2
end sub
def potential(x)       ! just a flat potential
   potential = 0
end def
sub check_normalization(p(,),m)    ! compute the normalization of psi*psi to be sure
   sum = 0                         ! that the algorithm is ok
   for i = 0 to m
      sum = sum + p(i,1)*p(i,1) + p(i,2)*p(i,2)
   next i
   print sum
end sub