!==========================================================================================
! Time Dependent Schrodinger Equation
!
! Implements Visscher's explicit scheme
!
!------------------------------------------------------------------------------------------
program     tdse_fd
 implicit none 
 integer   , parameter                   :: dp   = 8
 real  (dp), allocatable, dimension(:,:) :: repsi, impsi
 real  (dp), allocatable, dimension(  :) :: 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  (dp)                              :: k0,sigma,sigma2                 ! momentum, spread  of initial wave function
 real  (dp)                              :: xmax , tmax, x, t, dt, dx, a, amax, x0,norm, pexp, delpsi
 integer                                 :: Nx   , Nt
 integer                                 :: ix   , it
 real  (dp), parameter                   :: ZERO = 0.0_dp,ONE=1.0_dp,TWO=2.0_dp,HALF=0.5_dp,QUARTER=0.25_dp
 real  (dp), parameter                   :: tinypsi = 1.0e-20_dp            ! |psi| < tinypsi is set to 0
 real  (dp), parameter                   :: PI=atan2(ZERO,-ONE)
 integer                                 :: f_psi,f_rho                     ! file unit numbers
 integer                                 :: chk
 character(100)                          :: fmt_psi = "(I20,1000E30.15)"
!---------------------------------------------------------------
 print *,'# Enter x0, sigma, k0, dx, a, xmax, tmax, tskip: '
 read  *,         x0, sigma, k0, dx, a, xmax, tmax, tskip
!---------------------------------------------------------------
 Nx     = INT(TWO*xmax/ dx) + 1;
 dx     =     TWO*xmax/(Nx  - 1)            ! Recalculate dx, since requested dx might not fit exactly in [-xmax,xmax]
!---------------------------------------------------------------
 allocate(repsi(2,Nx),impsi(2,Nx),V(Nx),rho(Nx),stat=chk)
 if(chk /= 0) stop 'Allocation Failure'
!---------------------------------------------------------------
!Potential:
 V      = ZERO
 do ix  = 1, Nx
  x     = -xmax + (ix-1)*dx
  V(ix) = HALF  * x * x
! if(ABS(x) <= TWO  ) V(ix) = -12.0_dp
! if(ABS(x) <= HALF ) V(ix) =   6.0_dp
! V(ix) = ZERO
 end do
 call compute_amax                           ! Visscher stability condition
!if you wish to skip compute_amax: valid for HOC:
!amax  = ONE/(TWO + HALF * xmax * xmax * dx * dx) 
!---------------------------------------------------------------
 dt     = TWO   * a * dx * dx
 Nt     = INT(  tmax/ dt) + 1
 dt     =       tmax/(Nt-1)                  ! so that 0 <= t <= tmax
 a      = HALF  * dt/(dx*dx)                 ! Courant number
 sigma2 = sigma*sigma 
!---------------------------------------------------------------
 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       < 3 ) stop 'Nx too small'
 if( Nt/tskip < 2 ) stop 'Nt too small'
 open(newunit=f_rho,file='rho')
 open(newunit=f_psi,file='psi')
!---------------------------------------------------------------
!Initial condition:
 t               = ZERO
 rho             = ZERO
 repsi( : ,1 )   = ZERO ; impsi( : ,1 ) = ZERO
 repsi( : ,Nx)   = ZERO ; impsi( : ,Nx) = ZERO
 norm  = (TWO*PI*sigma2)**(-QUARTER)
 do ix = 2, Nx-1
  x    = -xmax   + (ix-1) *  dx
  pexp = QUARTER * (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) = ZERO
   impsi(old,ix) = ZERO
  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= ZERO
 do ix = 1, Nx
  x    = -xmax + (ix-1)*dx
  write(f_psi,fmt_psi) it, t, x, rho(ix), repsi(old,ix), impsi(old,ix)
 end do
 write (f_psi,*      ) ''
!---------------------------------------------------------------
!Time evolution:
!First evolve impsi by half time step:
 do ix = 2, Nx-1
  delpsi         =  a*(repsi(old,ix+1)-TWO*repsi(old,ix)+repsi(old,ix-1))-V(ix)*dt*repsi(old,ix)
  delpsi         = HALF*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)-TWO*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)-TWO*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   =  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,fmt_psi) it, t, x, rho(ix), repsi(new,ix), impsi(new,ix)
   end do
   write(f_psi,*) ''
  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)
!------------------------------------------------------------------------------------------
contains
!------------------------------------------------------------------------------------------
 subroutine     compute_amax

  real(8)               :: V1, V2, dt1, dt2
 
  V1  =  MAXVAL(V); if( V1 <= ZERO) V1 = ZERO
  V2  = -MINVAL(V); if( V2 <= ZERO) V2 = ZERO

  dt1 =  TWO/(V1 + (TWO/(dx*dx)) )
  dt2 =  HUGE(dt)

  if(V2>ZERO) dt2   = TWO/V2
  
  dt     = dt2
  if(dt1 < dt2) dt  = dt1
 
  amax= HALF * dt/(dx*dx)

 end subroutine compute_amax
!------------------------------------------------------------------------------------------
end program tdse_fd
!==========================================================================================
!  ---------------------------------------------------------------------
!  Copyright by Konstantinos N. Anagnostopoulos (2004-2022)
!  Physics Dept., National Technical University,
!  konstant@mail.ntua.gr, 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/>.
!  -----------------------------------------------------------------------
