scale_dft.F90

Go to the documentation of this file.

#include "scalelib.h"
! special 2D-DFT routines for low wavenumber spectral transform
module scale_dft
  use mpi
  use scale_precision
  use scale_prc, only: &
    prc_local_comm_world, &
    prc_myrank
  use scale_prc_cartesc, only: &
    prc_num_x, &
    prc_num_y, &
    prc_2drank
  use scale_comm_cartesc, only: &
    comm_datatype
  use scale_const, only: &
    pi => const_pi
  implicit none
 
  public :: dft_setup
  public :: dft_finalize
  public :: dft_g2g
  public :: dft_g2g_divfree
 
  integer, private :: IMAX, JMAX ! number of total grids for spectral transform
  integer, private :: IGS, JGS ! global index
  integer, private :: LMM, MMM ! maximum truncation wavenumber
 
  ! table for trigonometric function
  real(RP), private, allocatable :: table_x(:,:), table_y(:,:)
  real(RP), private, allocatable :: table_l(:), table_m(:)
  ! work array
  real(RP), private, allocatable :: work(:,:,:)
 
contains
 
  subroutine dft_setup(KA,KS,KE,IA,IS,IE,JA,JS,JE,LM,MM)
    implicit none
    integer,intent(in) :: ka, ia, ja
    integer,intent(in) :: ks, ke, is, ie, js, je
    integer,intent(in) :: lm, mm
    real(rp) :: x, y
    integer :: i, j, k, l, m
 
    imax = (ie-is+1)*prc_num_x
    jmax = (je-js+1)*prc_num_y
 
    igs = (ie-is+1)*prc_2drank(prc_myrank,1)+1
    jgs = (je-js+1)*prc_2drank(prc_myrank,2)+1
 
    lmm = lm
    mmm = mm
 
    allocate( table_x(is:ie,0:2*lm), table_y(js:je,0:2*mm) )
    allocate( table_l(0:2*lm), table_m(0:2*mm) )
    allocate( work(ka,0:2*lm,ja) )
 
    do i = is, ie
      x = 2*pi/imax*(i-is+igs)
      table_x(i,0) = 1
    enddo
 
    do l = 1, lm
      do i = is, ie
        x = 2*pi/imax*(i-is+igs)
        table_x(i,2*l-1) =  cos(l*x)
        table_x(i,2*l)   = -sin(l*x)
      enddo
    enddo
 
    do j = js, je
      y = 2*pi/jmax*(j-js+jgs)
      table_y(j,0) = 1
    enddo
 
    do m = 1, mm
      do j = js, je
        y = 2*pi/jmax*(j-js+jgs)
        table_y(j,2*m-1) =  cos(m*y)
        table_y(j,2*m)   = -sin(m*y)
      enddo
    enddo
 
    table_l(0) = 1
    do l = 1, lm
      table_l(2*l-1) = cos(pi*l/imax)
      table_l(2*l)   = sin(pi*l/imax)
    enddo
 
    table_m(0) = 1
    do m = 1, mm
      table_m(2*m-1) = cos(pi*m/jmax)
      table_m(2*m)   = sin(pi*m/jmax)
    enddo
 
    !$acc enter data copyin(table_x, table_y)
    !$acc enter data copyin(table_l, table_m)
    !$acc enter data create(work)
 
  end subroutine dft_setup
 
  subroutine dft_finalize
 
    !$acc exit data delete(table_x, table_y)
    !$acc exit data delete(table_l, table_m)
    !$acc exit data delete(work)
    deallocate( table_x, table_y )
    deallocate( table_l, table_m )
    deallocate( work )
 
    return
  end subroutine dft_finalize
 
  subroutine dft_g2s(KA,KS,KE,IA,IS,IE,JA,JS,JE,LM,MM,f,s)
    ! Grid to spectral transformation
    integer,intent(in) :: KA, IA, JA
    integer,intent(in) :: KS, KE, IS, IE, JS, JE
    integer,intent(in) :: LM, MM
    real(RP), intent(in) :: f(KA, IA, JA) ! x: full level, y: full level
    real(RP), intent(out) :: s(KA,0:2*LM,0:2*MM)
    real(RP) ::  work_s(KA,0:2*LM,0:2*MM)
    real(RP) :: c, tb
    integer :: i, j, k, l, m
    integer :: ierr
 
    !$acc data copyin(f) copyout(s) create(work_s)
 
    c = 1.0_rp/imax
    !$acc kernels
    do j = js, je
      do l = 0, 2*lm
        do k = ks, ke
          work(k,l,j) = 0
        enddo
        !$acc loop seq
        do i = is, ie
          tb = table_x(i,l)*c
          do k = ks, ke
            work(k,l,j) = work(k,l,j) + f(k,i,j)*tb
          enddo
        enddo
      enddo
    enddo
    !$acc end kernels
 
    !$acc kernels
    do m = 0, 2*mm
      do l = 0, 2*lm
        do k = ks, ke
          work_s(k,l,m) = 0
        enddo
      enddo
    enddo
    !$acc end kernels
 
    c = 1.0_rp/jmax
    !$acc kernels
    do m = 0, 2*mm
      do j = js, je
        tb = table_y(j,m)*c
        do l = 0, 2*lm
          do k = ks, ke
             !$acc atomic
            work_s(k,l,m) = work_s(k,l,m) + work(k,l,j)*tb
          enddo
        enddo
      enddo
    enddo
    !$acc end kernels
 
    call mpi_allreduce(work_s, s, ka*(2*lm+1)*(2*mm+1), comm_datatype, mpi_sum, prc_local_comm_world, ierr)
 
    !$acc end data
 
  end subroutine dft_g2s
 
  subroutine dft_s2g(KA,KS,KE,IA,IS,IE,JA,JS,JE,LM,MM,s,f)
    ! Grid to spectral transformation
    integer,intent(in) :: KA, IA, JA
    integer,intent(in) :: KS, KE, IS, IE, JS, JE
    integer,intent(in) :: LM, MM
    real(RP), intent(in) :: s(KA,0:2*LM,0:2*MM)
    real(RP), intent(out) :: f(KA, IA, JA) ! x: full level, y: full level
    real(RP) :: c, tb
    integer :: i, j, k, l, m
 
    !$acc data copyin(s) copyout(f)
 
    !$acc kernels
    do j = js, je
      do i = is, ie
        do k = ks, ke
          f(k,i,j) = 0
        enddo
      enddo
    enddo
    !$acc end kernels
 
    !$acc kernels
    do j = js, je
      do l = 0, 2*lm
        do k = ks, ke
          work(k,l,j) = 0
        enddo
      enddo
    enddo
    !$acc end kernels
 
    !$acc kernels
    do m = 0, 2*mm
      if( m == 0 ) then
        c = 1
      else
        c = 2
      endif
      do j = js, je
        tb = table_y(j,m)*c
        do l = 0, 2*lm
          do k = ks, ke
             !$acc atomic
             work(k,l,j) = work(k,l,j) + s(k,l,m)*tb
          enddo
        enddo
      enddo
    enddo
    !$acc end kernels
 
    !$acc kernels
    do j = js, je
      do i = is, ie
        do k = ks, ke
          f(k,i,j) = 0
        enddo
      enddo
      do l = 0, 2*lm
        if( l == 0 ) then
          c = 1
        else
          c = 2
        endif
        do i = is, ie
          tb = table_x(i,l)*c
          do k = ks, ke
             !$acc atomic
             f(k,i,j) = f(k,i,j) + work(k,l,j)*tb
          enddo
        enddo
      enddo
    enddo
    !$acc end kernels
 
    !$acc end data
 
  end subroutine dft_s2g
 
  ! Grid to Grid transformation
  subroutine dft_g2g(KA,KS,KE,IA,IS,IE,JA,JS,JE,LM,MM,f)
    integer,intent(in) :: ka, ia, ja
    integer,intent(in) :: ks, ke, is, ie, js, je
    integer,intent(in) :: lm, mm
    real(rp), intent(inout) :: f(ka, ia, ja) ! x,y: full or half level
    real(rp) :: s(ka,0:2*lm,0:2*mm)
 
    call dft_g2s(ka,ks,ke,ia,is,ie,ja,js,je,lm,mm,f,s)
    call dft_s2g(ka,ks,ke,ia,is,ie,ja,js,je,lm,mm,s,f)
 
  end subroutine dft_g2g
 
  ! Grid to Grid transformation of vector field with horizontal divergence filter
  subroutine dft_g2g_divfree(KA,KS,KE,IA,IS,IE,JA,JS,JE,LM,MM,u,v)
    integer,intent(in) :: ka, ia, ja
    integer,intent(in) :: ks, ke, is, ie, js, je
    integer,intent(in) :: lm, mm
    real(rp), intent(inout) :: u(ka, ia, ja) ! x: half level, y: full level
    real(rp), intent(inout) :: v(ka, ia, ja) ! x: full level, y: half level
    integer :: k, l, m
    real(rp) :: s1(ka,0:2*lm,0:2*mm)
    real(rp) :: s2(ka,0:2*lm,0:2*mm)
    real(rp) :: s3(ka,0:2*lm,0:2*mm)
    real(rp) :: a, b, fac
 
    !$acc data copy(u,v) create(s1,s2,s3)
 
    call dft_g2s(ka,ks,ke,ia,is,ie,ja,js,je,lm,mm,u,s1)
 
    ! phase shift
    !$acc kernels
    do m = 0, 2*mm
      do l = 1, lm
        do k = ks, ke
          a = s1(k,2*l-1,m)
          b = s1(k,2*l,m)
          s1(k,2*l-1,m) = a*table_l(2*l-1) + b*table_l(2*l)
          s1(k,2*l,m)   = b*table_l(2*l-1) - a*table_l(2*l)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    call dft_g2s(ka,ks,ke,ia,is,ie,ja,js,je,lm,mm,v,s2)
 
    ! phase shift
    !$acc kernels
    do m = 1, mm
      do l = 0, 2*lm
        do k = ks, ke
          a = s2(k,l,2*m-1)
          b = s2(k,l,2*m)
          s2(k,l,2*m-1) = a*table_m(2*m-1) + b*table_m(2*m)
          s2(k,l,2*m)   = b*table_m(2*m-1) - a*table_m(2*m)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! rotation
    !$acc kernels
    do m = 0, 2*mm
      do l = 0, 2*lm
        do k = ks, ke
          s3(k,l,m) = 0
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! dv/dx
    !$acc kernels
    do m = 0, 2*mm
      do l = 1, lm
        do k = ks, ke
          s3(k,2*l-1,m) = -l*s2(k,2*l,m)
          s3(k,2*l,m) = l*s2(k,2*l-1,m)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! + (- du/dy)
    !$acc kernels
    do m = 1, mm
      do l = 0, 2*lm
        do k = ks, ke
          s3(k,l,2*m-1) = s3(k,l,2*m-1) + m*s1(k,l,2*m)
          s3(k,l,2*m)   = s3(k,l,2*m)   - m*s1(k,l,2*m-1)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! minus inverse laplacian ( stream function on model plane )
    !$acc kernels
    do k = ks, ke
      s3(k,0,0) = 0
    enddo
    !$acc end kernels
 
    !$acc kernels
    do l = 1, lm
      fac = 1.0_rp/ (l*l)
      do k = ks, ke
        s3(k,2*l-1,0) = s3(k,2*l-1,0)*fac
        s3(k,2*l,0)   = s3(k,2*l,0)*fac
      enddo
    enddo
    !$acc end kernels
 
    !$acc kernels
    do m = 1, mm
      fac = 1.0_rp/ (m*m)
      do k = ks, ke
        s3(k,0,2*m-1) = s3(k,0,2*m-1)*fac
        s3(k,0,2*m)   = s3(k,0,2*m)*fac
      enddo
    enddo
    !$acc end kernels
 
    !$acc kernels
    do m = 1, mm
      do l = 1, lm
        fac = 1.0_rp/ (l*l+m*m)
        do k = ks, ke
          s3(k,2*l-1,2*m-1) = s3(k,2*l-1,2*m-1)*fac
          s3(k,2*l-1,2*m)   = s3(k,2*l-1,2*m)*fac
          s3(k,2*l,2*m-1)   = s3(k,2*l,2*m-1)*fac
          s3(k,2*l,2*m)     = s3(k,2*l,2*m)*fac
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! divergence free 2D vector
 
    ! dphi/dy
    !$acc kernels
    do l = 1, 2*lm
      do k = ks, ke
        s1(k,l,0) = 0
      enddo
    enddo
    !$acc end kernels
    !$acc kernels
    do m = 1, mm
      do l = 0, 2*lm
        do k = ks, ke
          s1(k,l,2*m-1) = -m*s3(k,l,2*m)
          s1(k,l,2*m)   =  m*s3(k,l,2*m-1)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! -dphi/dx
    !$acc kernels
    do m = 1, 2*mm
      do k = ks, ke
        s2(k,0,m) = 0
      enddo
    enddo
    !$acc end kernels
    !$acc kernels
    do m = 0, 2*mm
      do l = 1, lm
        do k = ks, ke
          s2(k,2*l-1,m) =  l*s3(k,2*l,m)
          s2(k,2*l,m)   = -l*s3(k,2*l-1,m)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    ! phase shift
    !$acc kernels
    do m = 0, 2*mm
      do l = 1, lm
        do k = ks, ke
          a = s1(k,2*l-1,m)
          b = s1(k,2*l,m)
          s1(k,2*l-1,m) = a*table_l(2*l-1) - b*table_l(2*l)
          s1(k,2*l,m)   = b*table_l(2*l-1) + a*table_l(2*l)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    call dft_s2g(ka,ks,ke,ia,is,ie,ja,js,je,lm,mm,s1,u)
 
    ! phase shift
    !$acc kernels
    do m = 1, mm
      do l = 0, 2*lm
        do k = ks, ke
          a = s2(k,l,2*m-1)
          b = s2(k,l,2*m)
          s2(k,l,2*m-1) = a*table_m(2*m-1) - b*table_m(2*m)
          s2(k,l,2*m) = b*table_m(2*m-1) + a*table_m(2*m)
        enddo
      enddo
    enddo
    !$acc end kernels
 
    call dft_s2g(ka,ks,ke,ia,is,ie,ja,js,je,lm,mm,s2,v)
 
    !$acc end data
 
  end subroutine dft_g2g_divfree
 
end module scale_dft