sll_m_poisson_2d_periodic_par.F90

Go to the documentation of this file.

 
!
!
module sll_m_poisson_2d_periodic_par
 
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#include "sll_memory.h"
#include "sll_assert.h"
#include "sll_working_precision.h"
 
   use sll_m_collective, only: &
      sll_t_collective_t, &
      sll_f_get_collective_rank, &
      sll_f_get_collective_size
 
   use sll_m_constants, only: &
      sll_p_pi
 
   use sll_m_fft, only: &
      sll_s_fft_exec_c2c_1d, &
      sll_p_fft_backward, &
      sll_s_fft_free, &
      sll_p_fft_forward, &
      sll_s_fft_init_c2c_1d, &
      sll_t_fft
 
   use sll_m_remapper, only: &
      sll_o_apply_remap_2d, &
      sll_o_compute_local_sizes, &
      sll_o_get_layout_collective, &
      sll_o_initialize_layout_with_distributed_array, &
      sll_t_layout_2d, &
      sll_o_local_to_global, &
      sll_f_new_layout_2d, &
      sll_o_new_remap_plan, &
      sll_t_remap_plan_2d_comp64, &
      sll_o_delete
 
   use sll_m_utilities, only: &
      sll_f_is_power_of_two
 
   implicit none
 
   public :: &
      sll_t_poisson_2d_periodic_par, &
      sll_s_poisson_2d_periodic_par_init, &
      sll_s_poisson_2d_periodic_par_init_alt, &
      sll_s_poisson_2d_periodic_par_solve, &
      sll_s_poisson_2d_periodic_par_solve_alt, &
      sll_s_poisson_2d_periodic_par_free
 
   private
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
   type sll_t_poisson_2d_periodic_par
      sll_int32                                 :: ncx    
      sll_int32                                 :: ncy    
      sll_real64                                :: lx     
      sll_real64                                :: ly     
      type(sll_t_fft)                           :: px
      type(sll_t_fft)                           :: py
      type(sll_t_fft)                           :: px_inv
      type(sll_t_fft)                           :: py_inv
      type(sll_t_layout_2d), pointer :: layout_seq_x1
      type(sll_t_layout_2d), pointer :: layout_seq_x2
      sll_int32                                 :: seq_x1_local_sz_x1 
      sll_int32                                 :: seq_x1_local_sz_x2 
      sll_int32                                 :: seq_x2_local_sz_x1 
      sll_int32                                 :: seq_x2_local_sz_x2 
      sll_comp64, dimension(:, :), pointer :: fft_x_array 
      sll_comp64, dimension(:, :), pointer :: fft_y_array 
      type(sll_t_remap_plan_2d_comp64), pointer :: rmp_xy
      type(sll_t_remap_plan_2d_comp64), pointer :: rmp_yx
      sll_comp64, pointer :: tmp_x(:)
      sll_comp64, pointer :: tmp_y(:)
 
   end type sll_t_poisson_2d_periodic_par
 
contains
 
   subroutine sll_s_poisson_2d_periodic_par_init(plan, &
                                                 start_layout, &
                                                 ncx, &
                                                 ncy, &
                                                 Lx, &
                                                 Ly)
 
      type(sll_t_poisson_2d_periodic_par) :: plan
      type(sll_t_layout_2d), pointer       :: start_layout
      sll_int32                            :: ncx          
      sll_int32                            :: ncy          
      sll_real64                           :: lx           
      sll_real64                           :: ly           
      sll_int64                            :: colsz ! collective size
      type(sll_t_collective_t), pointer    :: collective
 
      ! number of processors
      sll_int32                        :: nprocx1
      sll_int32                        :: nprocx2
      sll_int32                        :: ierr
      sll_int32                        :: loc_sz_x1
      sll_int32                        :: loc_sz_x2
 
      ! The collective to be used is the one that comes with the given layout.
      collective => sll_o_get_layout_collective(start_layout)
      colsz = int(sll_f_get_collective_size(collective), i64)
 
      if ((.not. sll_f_is_power_of_two(int(ncx, i64))) .and. &
          (.not. sll_f_is_power_of_two(int(ncy, i64)))) then
         print *, 'This test needs to run with numbers of cells which are', &
            'powers of 2.'
         print *, 'Exiting...'
         stop
      end if
 
      ! We use the number of cells since due to periodicity, the last point is
      ! not considered.
      plan%ncx = ncx
      plan%ncy = ncy
      plan%Lx = lx
      plan%Ly = ly
 
      ! Layout and local sizes for FFTs in x-direction
      plan%layout_seq_x1 => start_layout
      call sll_o_compute_local_sizes( &
         plan%layout_seq_x1, &
         loc_sz_x1, &
         loc_sz_x2)
 
      plan%seq_x1_local_sz_x1 = loc_sz_x1
      plan%seq_x1_local_sz_x2 = loc_sz_x2
 
      sll_allocate(plan%fft_x_array(loc_sz_x1, loc_sz_x2), ierr)
      sll_allocate(plan%tmp_x(ncx), ierr)
 
      ! For FFTs (in x-direction)
      call sll_s_fft_init_c2c_1d( &
         plan%px, &
         ncx, &
         plan%tmp_x, plan%tmp_x, &
         sll_p_fft_forward)!+FFT_NORMALIZE )
 
      call sll_s_fft_init_c2c_1d( &
         plan%px_inv, &
         ncx, &
         plan%tmp_x, plan%tmp_x, &
         sll_p_fft_backward) !+FFT_NORMALIZE )
 
      ! Layout and local sizes for FFTs in y-direction (x2)
      plan%layout_seq_x2 => sll_f_new_layout_2d(collective)
      nprocx1 = int(colsz, kind=4)
      nprocx2 = 1
 
      call sll_o_initialize_layout_with_distributed_array( &
         ncx + 1, &
         ncy + 1, &
         nprocx1, &
         nprocx2, &
         plan%layout_seq_x2)
 
      call sll_o_compute_local_sizes( &
         plan%layout_seq_x2, &
         loc_sz_x1, &
         loc_sz_x2)
 
      plan%seq_x2_local_sz_x1 = loc_sz_x1
      plan%seq_x2_local_sz_x2 = loc_sz_x2
      sll_allocate(plan%fft_y_array(loc_sz_x1, loc_sz_x2), ierr)
 
      sll_allocate(plan%tmp_y(ncy), ierr)
 
      ! For FFTs (in y-direction)
 
      call sll_s_fft_init_c2c_1d( &
         plan%py, &
         ncy, &
         plan%tmp_y, plan%tmp_y, &
         sll_p_fft_forward)! + FFT_NORMALIZE )
 
      call sll_s_fft_init_c2c_1d( &
         plan%py_inv, &
         ncy, &
         plan%tmp_y, plan%tmp_y, &
         sll_p_fft_backward)! + FFT_NORMALIZE )
 
      plan%rmp_xy => &
         sll_o_new_remap_plan(plan%layout_seq_x1, plan%layout_seq_x2, plan%fft_x_array)
      plan%rmp_yx => &
         sll_o_new_remap_plan(plan%layout_seq_x2, plan%layout_seq_x1, plan%fft_y_array)
   end subroutine sll_s_poisson_2d_periodic_par_init
 
   subroutine sll_s_poisson_2d_periodic_par_solve(plan, rho, phi)
      type(sll_t_poisson_2d_periodic_par) :: plan
      sll_real64, dimension(:, :)           :: rho      
      sll_real64, dimension(:, :)           :: phi      
      sll_int32                            :: ncx      
      sll_int32                            :: ncy      
      sll_int32                            :: npx_loc
      sll_int32                            :: npy_loc
      sll_int32                            :: i, j
      ! Reciprocals of domain lengths.
      sll_real64                           :: r_lx, r_ly
      sll_real64                           :: kx, ky
      sll_real64                           :: normalization
      type(sll_t_layout_2d), pointer       :: layout_x
      type(sll_t_layout_2d), pointer       :: layout_y
      sll_int32, dimension(1:2)            :: global
      sll_int32                            :: gi, gj
 
      ncx = plan%ncx
      ncy = plan%ncy
      r_lx = 1.0_f64/plan%Lx
      r_ly = 1.0_f64/plan%Ly
 
      !print*, 1.0_f64/plan%Lx, 1.0_f64/plan%Ly
      ! Get layouts to compute FFTs (in each direction)
      layout_x => plan%layout_seq_x1
      layout_y => plan%layout_seq_x2
      !call verify_argument_sizes_par(layout_x, rho, phi)
 
      ! FFTs in x-direction
      npx_loc = plan%seq_x1_local_sz_x1
      npy_loc = plan%seq_x1_local_sz_x2
 
      ! The input is handled internally as a complex array
      plan%fft_x_array = cmplx(rho, 0.0_f64, kind=f64)
 
      do j = 1, npy_loc
         call sll_s_fft_exec_c2c_1d(plan%px, plan%fft_x_array(1:ncx, j), plan%tmp_x)
         plan%fft_x_array(1:ncx, j) = plan%tmp_x
      end do
      ! FFTs in y-direction
      npx_loc = plan%seq_x2_local_sz_x1
      npy_loc = plan%seq_x2_local_sz_x2
 
      call sll_o_apply_remap_2d(plan%rmp_xy, plan%fft_x_array, plan%fft_y_array)
 
      do i = 1, npx_loc
         plan%tmp_y = plan%fft_y_array(i, 1:ncy)
         call sll_s_fft_exec_c2c_1d(plan%py, plan%tmp_y, plan%fft_y_array(i, 1:ncy))
      end do
 
      ! This should be inside the FFT plan...
      normalization = 1.0_f64/real(ncx*ncy, f64)
 
      ! Apply the kernel
      do j = 1, npy_loc - 1 ! last point was not transformed
         do i = 1, npx_loc - 1
            ! Make sure that the first mode is set to zero so that we get an
            ! answer with zero mean value. This step assumes that the (1,1) point
            ! will always be at the border of any splitting of the domains. This
            ! seems safe in this case.
 
            global = sll_o_local_to_global(layout_y, (/i, j/))
            gi = global(1)
            gj = global(2)
 
            if ((gi == 1) .and. (gj == 1)) then
               plan%fft_y_array(1, 1) = (0.0_f64, 0.0_f64)
            else
               kx = real(gi - 1, f64)
               ky = real(gj - 1, f64)
               ! Crucial step: While the results of the FFT are basically a
               ! list of the modes in the 0..n-1 sense, the algorithm that we
               ! are using uses a decomposition with negative frequencies,
               ! i.e.: -n/2..n/2-1. Therefore a step is needed in which we are
               ! effectively reinterpreting the FFT results in this light.
               ! More specifically, we take the second half of the transformed
               ! array and apply a modified factor, proper of a negative
               ! frequency.
               if (kx .ge. ncx/2) then
                  kx = kx - ncx
               end if
 
               if (ky .ge. ncy/2) then
                  ky = ky - ncy
               end if
 
               plan%fft_y_array(i, j) = -plan%fft_y_array(i, j)/ &
                                        (((kx*r_lx)**2 + (ky*r_ly)**2)*4.0_f64*sll_p_pi**2)
            end if
         end do
      end do
 
      ! Inverse FFTs in y-direction
      do i = 1, npx_loc
         plan%tmp_y = plan%fft_y_array(i, 1:ncy)
         call sll_s_fft_exec_c2c_1d(plan%py_inv, plan%tmp_y, plan%fft_y_array(i, :))
      end do
 
      ! Force the periodicity condition in the y-direction. CAN'T USE THE FFT
      ! INTERFACE SINCE THIS POINT FALLS OUTSIDE OF THE POINTS IN THE ARRAY
      ! TOUCHED BY THE FFT. This is another reason to permit not including the
      ! last point in the periodic cases...
 
      plan%fft_y_array(:, npy_loc) = plan%fft_y_array(:, 1)
 
      ! Prepare to take inverse FFTs in x-direction
      call sll_o_apply_remap_2d(plan%rmp_yx, plan%fft_y_array, plan%fft_x_array)
 
      npx_loc = plan%seq_x1_local_sz_x1
      npy_loc = plan%seq_x1_local_sz_x2
 
      do j = 1, npy_loc
         call sll_s_fft_exec_c2c_1d(plan%px_inv, plan%fft_x_array(1:ncx, j), plan%tmp_x)
         plan%fft_x_array(1:ncx, j) = plan%tmp_x
      end do
 
      ! Also ensure the periodicity in x
      plan%fft_x_array(npx_loc, :) = plan%fft_x_array(1, :)
 
      phi(1:npx_loc, 1:npy_loc) = real(plan%fft_x_array(1:npx_loc, 1:npy_loc), f64)
 
      phi = phi*normalization
 
   end subroutine sll_s_poisson_2d_periodic_par_solve
 
   subroutine sll_s_poisson_2d_periodic_par_init_alt( &
      plan, &
      start_layout, &
      ncx, &
      ncy, &
      Lx, &
      Ly)
 
      type(sll_t_poisson_2d_periodic_par) :: plan
      type(sll_t_layout_2d), pointer       :: start_layout
      sll_int32                            :: ncx          
      sll_int32                            :: ncy          
      sll_real64                           :: lx           
      sll_real64                           :: ly           
      sll_int64                            :: colsz        
      type(sll_t_collective_t), pointer    :: collective
      ! number of processors
      sll_int32                        :: nprocx1
      sll_int32                        :: nprocx2
      sll_int32                        :: ierr
      sll_int32                        :: loc_sz_x1
      sll_int32                        :: loc_sz_x2
 
      ! The collective to be used is the one that comes with the given layout.
      collective => sll_o_get_layout_collective(start_layout)
      colsz = int(sll_f_get_collective_size(collective), i64)
 
      sll_assert(sll_f_is_power_of_two(int(ncx, i64)))
      sll_assert(sll_f_is_power_of_two(int(ncy, i64)))
 
      ! We use the number of cells since due to periodicity, the last point is
      ! not considered. It should not even exist...
      plan%ncx = ncx
      plan%ncy = ncy
      plan%Lx = lx
      plan%Ly = ly
 
      ! Layout and local sizes for FFTs in x-direction
      plan%layout_seq_x1 => start_layout
      call sll_o_compute_local_sizes( &
         plan%layout_seq_x1, &
         loc_sz_x1, &
         loc_sz_x2)
 
      plan%seq_x1_local_sz_x1 = loc_sz_x1  ! ncx
      plan%seq_x1_local_sz_x2 = loc_sz_x2
 
      sll_allocate(plan%fft_x_array(loc_sz_x1, loc_sz_x2), ierr)
 
      allocate (plan%tmp_x(ncx))
      call sll_s_fft_init_c2c_1d( &
         plan%px, &
         ncx, &
         plan%tmp_x, &
         plan%tmp_x, &
         sll_p_fft_forward)!+FFT_NORMALIZE )
 
      call sll_s_fft_init_c2c_1d( &
         plan%px_inv, &
         ncx, &
         plan%tmp_x, &
         plan%tmp_x, &
         sll_p_fft_backward)!+FFT_NORMALIZE )
 
      ! Layout and local sizes for FFTs in y-direction (x2)
      plan%layout_seq_x2 => sll_f_new_layout_2d(collective)
      nprocx1 = int(colsz, kind=4)
      nprocx2 = 1
 
      call sll_o_initialize_layout_with_distributed_array( &
         ncx, &
         ncy, &
         nprocx1, &
         nprocx2, &
         plan%layout_seq_x2)
 
      call sll_o_compute_local_sizes( &
         plan%layout_seq_x2, &
         loc_sz_x1, &
         loc_sz_x2)
 
      plan%seq_x2_local_sz_x1 = loc_sz_x1
      plan%seq_x2_local_sz_x2 = ncy ! loc_sz_x2
      sll_allocate(plan%fft_y_array(loc_sz_x1, loc_sz_x2), ierr)
 
      ! For FFTs (in y-direction)
 
      allocate (plan%tmp_y(ncy))
      call sll_s_fft_init_c2c_1d( &
         plan%py, &
         ncy, &
         plan%tmp_y, &
         plan%tmp_y, &
         sll_p_fft_forward)! + FFT_NORMALIZE )
 
      call sll_s_fft_init_c2c_1d( &
         plan%py_inv, &
         ncy, &
         plan%tmp_y, &
         plan%tmp_y, &
         sll_p_fft_backward)! + FFT_NORMALIZE )
 
      plan%rmp_xy => &
         sll_o_new_remap_plan(plan%layout_seq_x1, plan%layout_seq_x2, plan%fft_x_array)
      plan%rmp_yx => &
         sll_o_new_remap_plan(plan%layout_seq_x2, plan%layout_seq_x1, plan%fft_y_array)
   end subroutine sll_s_poisson_2d_periodic_par_init_alt
 
   subroutine sll_s_poisson_2d_periodic_par_solve_alt(plan, rho, phi)
      type(sll_t_poisson_2d_periodic_par) :: plan
      sll_real64, dimension(:, :)           :: rho      
      sll_real64, dimension(:, :)           :: phi      
      sll_int32                            :: ncx      
      sll_int32                            :: ncy      
      sll_int32                            :: npx_loc
      sll_int32                            :: npy_loc
      sll_int32                            :: i, j
      ! Reciprocals of domain lengths.
      sll_real64                           :: r_lx, r_ly
      sll_real64                           :: kx, ky
      sll_real64                           :: normalization
      type(sll_t_layout_2d), pointer       :: layout_x
      type(sll_t_layout_2d), pointer       :: layout_y
      sll_int32, dimension(1:2)            :: global
      sll_int32                            :: gi, gj
 
      ncx = plan%ncx
      ncy = plan%ncy
      r_lx = 1.0_f64/plan%Lx
      r_ly = 1.0_f64/plan%Ly
 
      !print*, 1.0_f64/plan%Lx, 1.0_f64/plan%Ly
      ! Get layouts to compute FFTs (in each direction)
      layout_x => plan%layout_seq_x1
      layout_y => plan%layout_seq_x2
      !call verify_argument_sizes_par(layout_x, rho, phi)
 
      ! FFTs in x-direction
      npx_loc = plan%seq_x1_local_sz_x1
      npy_loc = plan%seq_x1_local_sz_x2
 
      ! The input is handled internally as a complex array
      plan%fft_x_array = cmplx(rho, 0.0_f64, kind=f64)
 
      do j = 1, npy_loc
         call sll_s_fft_exec_c2c_1d(plan%px, plan%fft_x_array(:, j), plan%tmp_x)
         plan%fft_x_array(:, j) = plan%tmp_x
      end do
      ! FFTs in y-direction
      npx_loc = plan%seq_x2_local_sz_x1
      npy_loc = plan%seq_x2_local_sz_x2
 
      call sll_o_apply_remap_2d(plan%rmp_xy, plan%fft_x_array, plan%fft_y_array)
 
      do i = 1, npx_loc
         plan%tmp_y = plan%fft_y_array(i, :)
         call sll_s_fft_exec_c2c_1d(plan%py, plan%tmp_y, plan%fft_y_array(i, :))
      end do
 
      ! This should be inside the FFT plan...
      normalization = 1.0_f64/real(ncx*ncy, f64)
 
      ! Apply the kernel
      do j = 1, npy_loc
         do i = 1, npx_loc
            ! Make sure that the first mode is set to zero so that we get an
            ! answer with zero mean value. This step assumes that the (1,1) point
            ! will always be at the border of any splitting of the domains. This
            ! seems safe in this case.
 
            global = sll_o_local_to_global(layout_y, (/i, j/))
            gi = global(1)
            gj = global(2)
 
            if ((gi == 1) .and. (gj == 1)) then
               plan%fft_y_array(1, 1) = (0.0_f64, 0.0_f64)
            else
               kx = real(gi - 1, f64)
               ky = real(gj - 1, f64)
               ! Crucial step: While the results of the FFT are basically a
               ! list of the modes in the 0..n-1 sense, the algorithm that we
               ! are using uses a decomposition with negative frequencies,
               ! i.e.: -n/2..n/2-1. Therefore a step is needed in which we are
               ! effectively reinterpreting the FFT results in this light.
               ! More specifically, we take the second half of the transformed
               ! array and apply a modified factor, proper of a negative
               ! frequency.
               if (kx .ge. ncx/2) then
                  kx = kx - ncx
               end if
 
               if (ky .ge. ncy/2) then
                  ky = ky - ncy
               end if
 
               plan%fft_y_array(i, j) = -plan%fft_y_array(i, j)/ &
                                        (((kx*r_lx)**2 + (ky*r_ly)**2)*4.0_f64*sll_p_pi**2)
            end if
         end do
      end do
 
      ! Inverse FFTs in y-direction
      do i = 1, npx_loc
         plan%tmp_y = plan%fft_y_array(i, :)
         call sll_s_fft_exec_c2c_1d(plan%py_inv, plan%tmp_y, plan%fft_y_array(i, :))
      end do
 
      ! Prepare to take inverse FFTs in x-direction
      call sll_o_apply_remap_2d(plan%rmp_yx, plan%fft_y_array, plan%fft_x_array)
 
      npx_loc = plan%seq_x1_local_sz_x1
      npy_loc = plan%seq_x1_local_sz_x2
 
      do j = 1, npy_loc
         call sll_s_fft_exec_c2c_1d(plan%px_inv, plan%fft_x_array(:, j), plan%tmp_x)
         plan%fft_x_array(:, j) = plan%tmp_x
      end do
 
      phi = real(plan%fft_x_array, f64)*normalization
 
   end subroutine sll_s_poisson_2d_periodic_par_solve_alt
 
   subroutine sll_s_poisson_2d_periodic_par_free(plan)
      type(sll_t_poisson_2d_periodic_par) :: plan
      sll_int32                            :: ierr
 
      call sll_s_fft_free(plan%px)
      call sll_s_fft_free(plan%py)
      call sll_s_fft_free(plan%px_inv)
      call sll_s_fft_free(plan%py_inv)
 
      call sll_o_delete(plan%layout_seq_x1)
      call sll_o_delete(plan%layout_seq_x2)
      sll_deallocate_array(plan%fft_x_array, ierr)
      sll_deallocate_array(plan%fft_y_array, ierr)
      call sll_o_delete(plan%rmp_xy)
      call sll_o_delete(plan%rmp_yx)
   end subroutine sll_s_poisson_2d_periodic_par_free
 
end module sll_m_poisson_2d_periodic_par