fft_obj_Transpose Subroutine

    impure subroutine fft_obj_Transpose(this,swap_1,swap_2,A,Ng,l_dims,l_offsets,A_T,Ng_T,l_dims_T,l_offsets_T)
      !> Builds the transposed array A_T by swapping directions (swap_1)
      ! and (swap_2) of input array A.
      ! Note that the MPI decomposition remains unchanged.
      ! E.g.:
      !  1) dir1=1 and dir2=2: field_T%cell(i,j,k) = field%cell(j,i,k)
      !  2) dir1=3 and dir2=2: field_T%cell(i,j,k) = field%cell(i,k,j)
      implicit none
      class(fft_obj),       intent(inout) :: this                              !! A FFT object
      integer,              intent(in)    :: swap_1                            !! Direction to swap
      integer,              intent(in)    :: swap_2                            !! Direction to swap
      real(wp),             intent(in)    :: A(:,:,:)                          !! Input array to transpose
      integer,              intent(in)    :: Ng(3)                             !! Global nbr of grid points in the original layout
      integer,              intent(in)    :: l_dims(3)                         !! Size of each subblock in the original layout
      integer,              intent(in)    :: l_offsets(3)                      !! Offsets of each subblock in the original layout
      real(wp),allocatable, intent(out)   :: A_T(:,:,:)                        !! Transposed array
      integer,              intent(out)   :: Ng_T(3)                           !! Global nbr of grid points in the transposed layout
      integer,              intent(out)   :: l_dims_T(3)                       !! Size of each subblock in the transposed layout
      integer,              intent(out)   :: l_offsets_T(3)                    !! Offsets of each subblock in the transposed layout
      ! Work variables
      integer :: t_idx(3), g_idx(3)
      integer :: rank
      integer :: target_rank
      integer :: perm(3)
      integer :: i,j,k
      integer :: ierr
      real(wp), allocatable :: s_buf(:), r_buf(:)

      ! Trivial case where the two directions to swap are the same
      if (swap_1.eq.swap_2) then
        ! Nothing to do
        allocate(A_T,mold=A)
        A_T = A

        return
      end if

      ! Determine data layout of transposed array
      call this%GetTransposeLayout(swap_1,swap_2,Ng,Ng_T,l_dims_T,l_offsets_T)

      ! Allocate A_T
      allocate(A_T(l_dims_T(1),l_dims_T(2),l_dims_T(3)), source = 0.0_WP)

      ! Count the number of elements to be sent to each rank from the current one
      this%s_counts = 0
      this%r_counts = 0
      do k=1,l_dims(3)
        do j=1,l_dims(2)
          do i=1,l_dims(1)
            ! Original indices
            g_idx = l_offsets + [i,j,k]

            ! Transposed indices
            t_idx         = g_idx
            t_idx(swap_1) = g_idx(swap_2)
            t_idx(swap_2) = g_idx(swap_1)

            ! Find rank where this element will be sent
            target_rank   = this%GetRankFromLayout(t_idx,Ng_T)
            this%s_counts(target_rank) = this%s_counts(target_rank) + 1
          end do
        end do
      end do

      ! Exchange counts so we know how much we are receiving
      call MPI_ALLTOALL(this%s_counts,1,this%parallel%INTEGER,this%r_counts,1,this%parallel%INTEGER,this%parallel%comm%g,ierr)

      ! Compute displacements
      this%s_displs(1) = 0
      this%r_displs(1) = 0
      do rank = 2,this%parallel%nproc
        this%s_displs(rank) = this%s_displs(rank-1) + this%s_counts(rank-1)
        this%r_displs(rank) = this%r_displs(rank-1) + this%r_counts(rank-1)
      end do

      ! Allocate send and receive buffers
      allocate(s_buf(sum(this%s_counts)))
      allocate(r_buf(sum(this%r_counts)))

      ! Initialize buffer pointers
      this%s_ptr = this%s_displs
      this%r_ptr = this%r_displs

      ! Pack data
      do k=1,l_dims(3)
        do j=1,l_dims(2)
          do i=1,l_dims(1)
            ! Global indices
            g_idx = l_offsets + [i,j,k]

            ! Transposed indices
            t_idx = g_idx
            t_idx(swap_1) = g_idx(swap_2)
            t_idx(swap_2) = g_idx(swap_1)

            ! Find rank where this element will be sent
            target_rank   = this%GetRankFromLayout(t_idx,Ng_T)
            s_buf(this%s_ptr(target_rank) + 1) = A(i,j,k)
            this%s_ptr(target_rank) = this%s_ptr(target_rank) + 1
          end do
        end do
      end do

      ! Global communication
      call MPI_ALLTOALLV(s_buf, this%s_counts, this%s_displs, this%parallel%REAL_WP, &
                         r_buf, this%r_counts, this%r_displs, this%parallel%REAL_WP, this%parallel%comm%g,ierr)

      ! Unpack data
      perm=[1,2,3]
      perm(swap_1) = swap_2
      perm(swap_2) = swap_1
      do k = 1,l_dims_T(perm(3))
        do j = 1,l_dims_T(perm(2))
          do i = 1,l_dims_T(perm(1))

            g_idx = [i,j,k]
            t_idx(1) = g_idx(perm(1))
            t_idx(2) = g_idx(perm(2))
            t_idx(3) = g_idx(perm(3))
            g_idx = [l_offsets_T(perm(1)),l_offsets_T(perm(2)),l_offsets_T(perm(3))] + [i,j,k]

            ! Figure out from which rank this entry was recieved and add it
            rank = this%GetRankFromLayout(g_idx, Ng)
            A_T(t_idx(1), t_idx(2), t_idx(3)) = r_buf(this%r_ptr(rank) + 1)
            this%r_ptr(rank) = this%r_ptr(rank) + 1
          end do
        end do
      end do

      deallocate(s_buf)
      deallocate(r_buf)

      return
    end subroutine fft_obj_Transpose