Program HeatedPlateCrankNicholsonDenseLU

  ! 2D Heat Equation
  ! Finite difference
  ! Crank-Nicholson
  ! dense LU i.e. dgesv     

  implicit none

  integer                 :: ii, kk, iBlock, iRow              ! counters
  real(kind=8), parameter :: C  = 1.0d+0                       ! thermal diffusivity (m^2/s) (gold is 1.27E-4)
  real(kind=8), parameter :: x  = 1.0d+0                       ! side of square (m)
  real(kind=8), parameter :: t  = 1.0d+0                       ! simulation time (s)
  real(kind=8), parameter :: dx = 0.25d+0                      ! x grid spacing (0.01=1cm)
  integer,      parameter :: m = (x/dx)                        ! side of square domain
  real(kind=8), parameter :: dy = dx                           ! square grid cells
  real(kind=8), parameter :: dt = m*(dx**2 * dx**2)/(2.0d+0*C*(dx**2 + dx**2))  ! can take bigger time-steps, say m*CFL
  real(kind=8), parameter :: a1 = 1.0d+0 + 2.0d+0*C*dt/(dx*dx) ! for use populating the CN A matrix
  real(kind=8), parameter :: a2 = 1.0d+0 - 2.0d+0*C*dt/(dx*dx) ! for use populating the CN B matrix
  real(kind=8), parameter :: c1 = -C*dt/(2.0d+0*(dx*dx))       ! for use populating the CN A matrix
  real(kind=8), parameter :: c2 = C*dt/(2.0d+0*(dx*dx))        ! for use populating the CN B matrix
  integer, parameter    :: KL   = m   ! number of sub-diagonals within band of matrices A and B
  integer, parameter    :: KU   = m   ! number of super-diagonals within band of matrices A and B
  integer,      parameter :: INCX  = 1
  integer,      parameter :: INCY  = 1
  integer,      parameter :: NRHS  = 1
  integer                 :: INFO
  character,    parameter :: TRANS = 'N'
  real(kind=8), parameter :: ALPHA = 1.0d+0
  real(kind=8), parameter :: BETA  = 0.0d+0
  real(kind=8), parameter :: epsilon = 1.0E-7                  ! convergence threshold
  logical                 :: converged = .false.               ! has the simulation converged?
  integer                 :: lastiter = 0                      ! last iteration if converged
  integer,      parameter :: N = m*m                           ! no. of cells in domain
  integer,      parameter :: LDA  = 2 * KL + KU + 1
  integer,      parameter :: LDB  = KL + KU + 1
  integer                 :: AROW   ! KL + KU + 1 + row(A) - col(A)
  integer                 :: BROW   ! KU + 1 + row(A) - col(A)
  integer                 :: start_count, end_count, count_rate, count_max ! for timing
  integer, allocatable, dimension(:) :: IPIV 
  real(kind=8), allocatable, dimension(:,:)         :: A
  real(kind=8), allocatable, dimension(:,:)         :: B
  real(kind=8), target, allocatable, dimension(:,:) :: U1,U2
  real(kind=8), pointer, dimension(:,:)             :: UK, UKP1, temp  ! for fast storage re-use

  !! allocate memory
  allocate(A(LDA,N))
  allocate(B(LDB,N))
  allocate(U1(m,m))
  allocate(U2(m,m))
  allocate(IPIV(N))

  !! initialise 
  ! pointers
  UK => U1
  UKP1 => U2
  ! values
  A    = 0.0d+0
  B    = 0.0d+0
  UK   = 0.0d+0
  UKP1 = 0.0d+0
  ! and boundary conditions
  ! top row
  UK(m,:)   = 100.0d+0
  UKP1(m,:) = 100.0d+0
  ! first & last columns
  UK(:,1)   = 100.0d+0
  UK(:,m)   = 100.0d+0
  UKP1(:,1) = 100.0d+0
  UKP1(:,m) = 100.0d+0

  !! populate A and B matrices
  do iBlock=1,m
     do ii=1,m
        iRow = (iBlock - 1)*m + ii
        if ((iBlock == 1) .or. (iBlock == m)) then ! first of last block
           AROW = KL + KU + 1 + iRow - iRow
           A(AROW,iRow) = 1.0d+0
           BROW = KU + 1 + iRow - iRow
           B(BROW,iRow) = 1.0d+0
        else                                       ! other blocks
           if ((ii == 1) .or. (ii == m)) then      ! first or last row in block
              AROW = KL + KU + 1 + iRow - iRow
              A(AROW,iRow) = 1.0d+0
              BROW = KU + 1 + iRow - iRow
              B(BROW,iRow) = 1.0d+0
           else                                    ! full 5-pt stencil
              AROW = KL + KU + 1 + iRow - iRow
              A(AROW,iRow)   = a1
              AROW = KL + KU + 1 + iRow - (iRow+1)
              A(AROW,iRow+1) = c1
              AROW = KL + KU + 1 + iRow - (iRow-1)
              A(AROW,iRow-1) = c1
              AROW = KL + KU + 1 + iRow - (iRow+m)
              A(AROW,iRow+m) = c1
              AROW = KL + KU + 1 + iRow - (iRow-m)
              A(AROW,iRow-m) = c1
              BROW = KU + 1 + iRow - iRow
              B(BROW,iRow)   = a2
              BROW = KU + 1 + iRow - (iRow+1)
              B(BROW,iRow+1) = c2
              BROW = KU + 1 + iRow - (iRow-1)
              B(BROW,iRow-1) = c2
              BROW = KU + 1 + iRow - (iRow+m)
              B(BROW,iRow+m) = c2
              BROW = KU + 1 + iRow - (iRow-m)
              B(BROW,iRow-m) = c2
           end if
        end if        
     end do
  end do

  call system_clock(start_count,count_rate,count_max) ! start timing

  !! want to solve Ax = By
  ! B*UK = UKP1
  ! NB specify the size that the unbanded matrix would be,
  ! then the size of the banded storage used to hold it
  call dgbmv(TRANS, N, N, KL, KU, ALPHA, B, LDB, UK, INCX, BETA, UKP1, INCY)

  ! soln is stored in place in UKP1
  call dgbsv(N, KL, KU, NRHS, A, LDA, IPIV, UKP1, N, INFO)

  if (INFO .ne. 0) then
     write(*,*), 'ERROR calling dgesv'
     stop
  endif

  do kk=2,int(t/dt)

     ! cycle
     temp => UK
     UK => UKP1
     UKP1 => temp

     call dgbmv(TRANS, N, N, KL, KU, ALPHA, B, LDB, UK, INCX, BETA, UKP1, INCY)

     ! Now have LU decomposition stored in A, so can call dgetrs instead of dgesv
     call dgbtrs('N', N, KL, KU, NRHS, A, LDA, IPIV, UKP1, N, INFO)
     
     if (INFO .ne. 0) then
        write(*,*), 'ERROR calling dgetrs'
        stop
     endif

#ifdef DEBUG
     ! test for convergence
     if (maxval(ABS(UKP1-UK))/dt <= epsilon) then
        converged = .true.
        lastiter = kk
        exit
     end if
#endif
     
  end do
  call system_clock(end_count) ! end timing

#ifdef DEBUG
  !! write out result--row-by-row
  do ii=M,1,-1
     write(*,*) UKP1(ii,:)
  end do
#endif

  !! report
  if (count_rate .ne. 1000) then
     write(*,*), 'ERROR timing is not in milliseconds'
     stop
  endif
  write(*,*) 'solving for', N, 'cells'
  write(*,*) 'timestep (s):', dt, '(', int(1.0/dt), 'iters/simulation second)'
  write(*,*) 'maxiter:', int(t/dt)
  if (converged) then
     write(*,*) 'Solutution converged to', epsilon, 'at iteration', lastiter
  end if
  write(*,*) 'solution took: ', real(end_count-start_count), '(ms)'

end Program HeatedPlateCrankNicholsonDenseLU