fft_fftw_mod.F90

Go to the documentation of this file.

module fft
  use, intrinsic :: iso_c_binding
  use fftw3, only : fftw_r2hc, fftw_hc2r, fftw_redft00, fftw_redft10, fftw_redft01
  implicit none

  public :: fft_r2r, transpose_grid, wavenumber
  public :: w_forward, w_backward, p_forward, p_backward 

  private

  integer, parameter :: nplans = 10
  integer :: n_r2r_plans = 0
  type(c_ptr) :: r2r_plans(nplans)
  integer :: r2r_plan_lengths(nplans)
  integer :: r2r_plan_nums(nplans)
  integer(C_INT) :: r2r_plan_kinds(nplans)

  integer, parameter :: w_forward = fftw_redft10
  integer, parameter :: w_backward = fftw_redft01
  integer, parameter :: p_forward = fftw_r2hc
  integer, parameter :: p_backward = fftw_hc2r

contains

subroutine fft_r2r(u, length, num, kind, rescale)
  use kinds, only : dp
  use fftw3, only : fftw_exhaustive, fftw_estimate
  use fftw3, only : fftw_plan_many_r2r
  use fftw3, only : fftw_execute_r2r
  use parallel, only : nid

  real(DP), intent(inout) :: u(:,:,:) !>
  integer,  intent(in)    :: length  !>
  integer,  intent(in)    :: num     !>
  integer(C_INT)          :: kind    !>
  real(DP), intent(inout) :: rescale !>

  integer :: i, plan_idx
  real(DP), allocatable :: proxy(:)

  ! Look for an existing plan
  plan_idx = -1
  do i = 1, n_r2r_plans
    if ( &
      length == r2r_plan_lengths(i) .and. &
      num    == r2r_plan_nums(i)    .and. &
      kind   == r2r_plan_kinds(i) &
       ) then
      plan_idx = i
      exit
    endif
  enddo

  ! If there's no plan, make a new one
  if (plan_idx < 0) then
    if (n_r2r_plans == nplans) then
      write(*,*) "Ran out of plans!"
    endif

    n_r2r_plans = n_r2r_plans + 1
    plan_idx = n_r2r_plans
    allocate(proxy(num*length))
    r2r_plans(plan_idx) = fftw_plan_many_r2r(1, &
                            (/length/), num, &
                            proxy, (/length/), 1, length, &
                            proxy, (/length/), 1, length, &
                            (/kind/), fftw_exhaustive) 
    deallocate(proxy)
    r2r_plan_lengths(plan_idx) = length
    r2r_plan_nums(plan_idx)    = num
    r2r_plan_kinds(plan_idx)   = kind
  endif

  ! execute the plan
  call fftw_execute_r2r(r2r_plans(plan_idx), u, u)

  ! rescale the normalization factor
  if (kind == fftw_redft00) then
    rescale = rescale * sqrt(2.*real(length-1, kind=dp))
  else if (kind == w_forward .or. kind == w_backward) then
    rescale = rescale * sqrt(2.*real(length, kind=dp))
  else if (kind == fftw_r2hc .or. kind == fftw_hc2r) then
    rescale = rescale * sqrt(real(length, kind=dp))
  else
    write(*,*) "Don't know how to deal with FFT kind", kind
  endif

end subroutine fft_r2r

real(DP) function wavenumber(i, N, L, kind)
  use kinds, only : dp
  implicit none
  integer,  intent(in) :: i    !>
  integer,  intent(in) :: N    !>
  real(DP), intent(in) :: L    !>
  integer,  intent(in) :: kind !>

  real(DP), parameter :: pi = 4.*atan(1.)

  if (kind == p_forward) then
    if (i <= n / 2) then
      wavenumber = 2*pi*i/(l)
    else
      wavenumber = 2*pi*(n - i)/(l)
    endif
  else if (kind == w_forward) then
    wavenumber = pi*i/(l)
  else
    write(*,*) "Don't know how to deal with FFT kind", kind
  endif

end function wavenumber

subroutine transpose_grid(grid, grid_t, shape_x, idx, idx_t, comm)
  use kinds, only : dp
  use fftw3, only : fftw_exhaustive, fftw_estimate
  use parallel, only : nid, nekreal

  real(DP), intent(inout) :: grid(0:,0:,0:)
  real(DP), intent(out)   :: grid_t(0:,0:,0:)
  integer,  intent(in)    :: shape_x(3)
  integer, intent(in)     :: idx
  integer, intent(in)     :: idx_t
  integer, intent(in)     :: comm

  real(DP), allocatable :: tmp(:,:)
  real(DP), allocatable :: tmp_t(:,:)
  integer :: block0, block1, num
  integer :: i, j, k, ierr, comm_size


  if (size(grid) < 1) return

  call mpi_comm_size(comm, comm_size, ierr)
 
  if (idx == 1 .or. idx_t == 1) then
    block0 = size(grid,2)
    block1 = size(grid,1) / comm_size
    num    = size(grid,3) 
  else if (idx == 3 .or. idx_t == 3) then
    block0 = size(grid,3)
    block1 = size(grid,1) / comm_size
    num    = size(grid,2) 
  endif
  allocate(tmp(0:block0-1,   0:size(grid,1)-1))
  allocate(tmp_t(0:block0-1, 0:size(grid,1)-1))

  if (idx == 1 .or. idx_t == 1) then
    do i = 0, num - 1
      tmp = transpose(grid(:,:,i))
      call mpi_alltoall(tmp,   block0*block1, nekreal, &
                        tmp_t, block0*block1, nekreal, comm, ierr)
      if (ierr /= 0) write(*,*) "alltoall errored", ierr, nid
      do j = 0, size(grid_t,1) - 1
        do k = 0, size(grid_t,2) - 1
          grid_t(j,k,i) = tmp_t( mod(j,int(block0)), (j / block0) * block1 + k )
        enddo
      enddo
    enddo
  else if (idx == 3 .or. idx_t == 3) then
    do i = 0, num - 1
      tmp = transpose(grid(:,i,:))
      call mpi_alltoall(tmp,   block0*block1, nekreal, &
                        tmp_t, block0*block1, nekreal, comm, ierr)
      if (ierr /= 0) write(*,*) "alltoall errored", ierr, nid
      do j = 0, size(grid_t,1) - 1
        do k = 0, size(grid_t,3) - 1
          grid_t(j,i,k) = tmp_t( mod(j,int(block0)), (j / block0) * block1 + k )
        enddo
      enddo
    enddo
  else
    write(*,*) "Something went wrong in transpose", nid
  endif

end subroutine transpose_grid

end module