!==========================================================================================
! Time Dependent Schrodinger Equation
!
! Implements Visscher's explicit scheme
!
!------------------------------------------------------------------------------------------
program     tdse_fd
 implicit none
 integer   , parameter      :: P = 100000
 real   (8), dimension(2,P) :: repsi, impsi
 complex(8), dimension(  P) :: psi
 real   (8), dimension(  P) :: V    , rho                      ! potential, rho=|psi|^2
 integer   , parameter      :: old=1, new=2                    ! indices in repsi, impsi
 integer                    :: tskip                           ! every how many steps to measure
 real   (8)                 :: k0,sigma,sigma2                 ! momentum, spread  of initial wave function
 real   (8)                 :: xmax , tmax, x, t, dt, dx, a, amax, x0,norm, pexp, delpsi,xmin
 real   (8)                 :: avX,avP,avX2,avP2,avE,DX2,DP2
 integer                    :: Nx   , Nt
 integer                    :: ix   , it
 real   (8), parameter      :: tinypsi = 1.0D-20               ! |psi| < tinypsi is set to 0
 real   (8), parameter      :: PI=atan2(1.0D0,0.0D0)*2.0D0
 integer, parameter         :: f_psi=11,f_rho=12,f_obs=13      ! file unit numbers
!---------------------------------------------------------------
 print *,'# Enter x0, sigma, k0, dx, a, xmax, tmax, tskip: '
 read  *,         x0, sigma, k0, dx, a, xmax, tmax, tskip
!---------------------------------------------------------------
 Nx     = INT(2.0D0*xmax/dx) + 1;
 dx     =     2.0D0*xmax/(Nx - 1)            ! Recalculate dx, since requested dx might not fit exactly in [-xmax,xmax]
!---------------------------------------------------------------
!Potential:
 V      = 0.0D0
 do ix  = 1, Nx
  x     = -xmax + (ix-1)*dx
  V(ix) = 0.5D0 * x * x
! if(ABS(x) <= 2.0D0 ) V(ix) = -12.0D0
! if(ABS(x) <= 0.5D0 ) V(ix) =   6.0D0
! V(ix) = 0.0D0
 end do
 call compute_amax(V,Nx,dx,amax)             ! Visscher stability condition
!---------------------------------------------------------------
 dt     = 2.0D0 * a * dx * dx
 Nt     = INT(  tmax/dt) + 1
 dt     =       tmax/(Nt-1)                  ! so that 0 <= t <= tmax
 a      = 0.5D0 * dt/(dx*dx)                 ! Courant number
 sigma2 = sigma*sigma                        ! 
 xmin   = -xmax
!---------------------------------------------------------------
 print *,'# ----------------------------------------------------'
 print *,'#    x0, sigma,  k0:'
 print *,'# ', x0, sigma,  k0
 print *,'#    xmax, tmax, tskip:'
 print *,'# ', xmax, tmax, tskip
 print *,'#    Nx, Nt, dx, dt:'
 print *,'# ', Nx, Nt, dx, dt
 print *,'#    a, amax:'
 print *,'# ', a, amax
 if( a>amax ) print *,'# WARNING: a>amax'
 print *,'# ----------------------------------------------------'
 if( Nx       > P ) stop 'Nx > P'
 if( Nx       < 3 ) stop 'Nx too small'
 if( Nt/tskip < 2 ) stop 'Nt too small'
 open(unit=f_rho,file='rho')
 open(unit=f_psi,file='psi')
 open(unit=f_obs,file='obs')
 write(f_obs,'(A                )') '#     Nx, Nt, x0, sigma, k0, xmax, tmax, tskip, dx, dt, a:'
 write(f_obs,'(A,2I12,1000E26.16)') '# ',  Nx, Nt, x0, sigma, k0, xmax, tmax, tskip, dx, dt, a
!---------------------------------------------------------------
!Initial condition:
 t = 0.0D0;
 rho           = 0.0D0
 repsi( : ,1 ) = 0.0D0; impsi( : ,1 ) = 0.0D0
 repsi( : ,Nx) = 0.0D0; impsi( : ,Nx) = 0.0D0
 norm  = (2.0D0*PI*sigma2)**(-0.25D0)
 do ix = 2, Nx-1
  x = -xmax + (ix-1)*dx
  pexp = 0.25D0 * (x-x0) * (x-x0) / sigma2
  if( pexp < -log(tinypsi))then !avoid very small numbers
   repsi(old,ix) = norm*exp( -pexp ) * cos( k0 * (x-x0) )
   impsi(old,ix) = norm*exp( -pexp ) * sin( k0 * (x-x0) )
  else
   repsi(old,ix) = 0.0d0
   impsi(old,ix) = 0.0d0
  end if
  rho  (     ix) = repsi(old,ix)*repsi(old,ix)+impsi(old,ix)*impsi(old,ix)
 end do
 norm =  SUM(rho)*dx  !  ∫ rho(x) dx
!---------------------------------------------------------------
!Output initial state:
 it =0; t= 0.0D0
 do ix = 1, Nx
  x = -xmax + (ix-1)*dx
  write(f_psi,101) it, t, x, rho(ix), repsi(old,ix), impsi(old,ix)
 end do
 write (f_psi,*) '';
 psi=DCMPLX(repsi(old,:),impsi(old,:))
 call observables(psi,V,Nx,xmin,dx, avX,avP,avX2,avP2,avE,DX2,DP2)
 write (f_obs,101) it, t,           avX,avP,avX2,avP2,avE,DX2,DP2
!---------------------------------------------------------------
!Time evolution:
!First evolve impsi by half time step:
 do ix = 2, Nx-1
  delpsi         =  a*(repsi(old,ix+1)-2.0D0*repsi(old,ix)+repsi(old,ix-1))-V(ix)*dt*repsi(old,ix)
  delpsi         = 0.5D0*delpsi
  if(abs(delpsi) > tinypsi)  impsi(old,ix)  = impsi(old,ix) + delpsi
 end do
!---------------------------------------------------------------
!Begin loop on time evolution:
 do it = 1, Nt
  t = it*dt
  !update repsi:
  do ix = 2, Nx-1
   delpsi        = -a*(impsi(old,ix+1)-2.0D0*impsi(old,ix)+impsi(old,ix-1))+V(ix)*dt*impsi(old,ix)
   if(abs(delpsi)> tinypsi) repsi(new,ix) = repsi(old,ix) + delpsi
   !We square the numbers repsi,impsi, therefore ABS(repsi*repsi),ABS(impsi*impsi) < TINY(delpsi)
   rho  (    ix) = repsi(old,ix)*repsi(new,ix)+impsi(old,ix)*impsi(old,ix)
  end do
  !update impsi:
  do ix = 2, Nx-1
   delpsi        =  a*(repsi(new,ix+1)-2.0D0*repsi(new,ix)+repsi(new,ix-1))-V(ix)*dt*repsi(new,ix)
   if(abs(delpsi)> tinypsi) impsi(new,ix) = impsi(old,ix) + delpsi
  end do
  !---------------------------------------------------------------
  !print norm of psi:  ∫ rho(x) dx
  norm   =  sqrt(SUM(rho)*dx)
  write(f_rho,*) it, t, norm
  !print psi:
  if(MOD(it,tskip) == 0)then
   do ix = 1, Nx
    x    = -xmax + (ix-1)*dx
    write(f_psi,101) it, t, x, rho(ix), repsi(new,ix), impsi(new,ix)
   end do
   write(f_psi,*) ''
  !---------------------------------------------------------------
  !compute observables:
   psi=DCMPLX(repsi(new,:),0.5D0*(impsi(new,:)+impsi(old,:)))
   call observables(psi,V,Nx,xmin,dx, avX,avP,avX2,avP2,avE,DX2,DP2)
   write (f_obs,101) it, t,           avX,avP,avX2,avP2,avE,DX2,DP2
  end if
  !---------------------------------------------------------------
  !interchange new<->old values:
  repsi(old,:) = repsi(new,:)
  impsi(old,:) = impsi(new,:)
 end do ! do it = 1, Nt    end time evolution
!---------------------------------------------------------------
 close(f_rho)
 close(f_psi)
101 FORMAT(I20,1000E30.15)
end program tdse_fd
!==========================================================================================
!------------------------------------------------------------------------------------------
subroutine     compute_amax(V,Nx,dx,a)
 implicit none
 integer               :: Nx
 real(8),dimension(Nx) :: V
 real(8)               :: dx,a
 real(8)               :: V1, V2, dt1, dt2,dt
 
 V1 =  MAXVAL(V(1:Nx)); if( V1 <= 0.0D0) V1 = 0.0D0
 V2 = -MINVAL(V(1:Nx)); if( V2 <= 0.0D0) V2 = 0.0D0

 dt1 = 2.0D0/(V1 + (2.0D0/(dx*dx)) )
 dt2 = HUGE(dt)
 if(V2>0.0D0) dt2   = 2.0D0/V2

 dt  = dt2
 if(dt1<dt2) dt = dt1

 a   = 0.5D0 * dt/(dx*dx)

end subroutine compute_amax
!==========================================================================================
!  ---------------------------------------------------------------------
!  Copyright by Konstantinos N. Anagnostopoulos (2004-2020)
!  Physics Dept., National Technical University,
!  konstant@mail.ntua.gr, www.physics.ntua.gr/~konstant
!  
!  This program is free software: you can redistribute it and/or modify
!  it under the terms of the GNU General Public License as published by
!  the Free Software Foundation, version 3 of the License.
!  
!  This program is distributed in the hope that it will be useful, but
!  WITHOUT ANY WARRANTY; without even the implied warranty of
!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
!  General Public License for more details.
!  
!  You should have received a copy of the GNU General Public Liense along
!  with this program.  If not, see <http://www.gnu.org/licenses/>.
!  -----------------------------------------------------------------------
