sll_m_poisson_2d_polar_par.F90

Go to the documentation of this file.

!**************************************************************
!  Copyright INRIA
!  Authors :
!     CALVI project team
!
!  This code SeLaLib (for Semi-Lagrangian-Library)
!  is a parallel library for simulating the plasma turbulence
!  in a tokamak.
!
!  This software is governed by the CeCILL-B license
!  under French law and abiding by the rules of distribution
!  of free software.  You can  use, modify and redistribute
!  the software under the terms of the CeCILL-B license as
!  circulated by CEA, CNRS and INRIA at the following URL
!  "http://www.cecill.info".
!**************************************************************
 
 
module sll_m_poisson_2d_polar_par
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#include "sll_assert.h"
#include "sll_errors.h"
#include "sll_working_precision.h"
 
   use sll_m_boundary_condition_descriptors, only: &
      sll_p_dirichlet, &
      sll_p_neumann, &
      sll_p_neumann_mode_0, &
      sll_p_polar_origin
 
   use sll_m_collective, only: &
      sll_f_get_collective_size, &
      sll_s_halt_collective, &
      sll_v_world_collective
 
   use sll_m_fft, only: &
      sll_t_fft, &
      sll_p_fft_forward, &
      sll_p_fft_backward, &
      sll_f_fft_allocate_aligned_real, &
      sll_s_fft_deallocate_aligned_real, &
      sll_s_fft_init_r2r_1d, &
      sll_s_fft_exec_r2r_1d, &
      sll_s_fft_get_k_list_r2r_1d, &
      sll_s_fft_free
 
   use sll_m_remapper, only: &
      sll_t_layout_2d, &
      sll_o_compute_local_sizes, &
      sll_o_get_layout_global_size_i, &
      sll_o_get_layout_global_size_j, &
      sll_o_local_to_global, &
      sll_t_remap_plan_2d_real64, &
      sll_o_new_remap_plan, &
      sll_o_apply_remap_2d
 
   use sll_m_tridiagonal, only: &
      sll_s_setup_cyclic_tridiag, &
      sll_o_solve_cyclic_tridiag
 
   use sll_m_utilities, only: &
      sll_s_new_array_linspace
 
   implicit none
 
   public :: &
      sll_t_poisson_2d_polar_par, &
      sll_s_poisson_2d_polar_par_init, &
      sll_s_poisson_2d_polar_par_solve, &
      sll_s_poisson_2d_polar_par_free
 
   private
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
   type sll_t_poisson_2d_polar_par
 
      real(f64)              :: rmin
      real(f64)              :: rmax
      integer(i32)              :: nr
      integer(i32)              :: ntheta
      integer(i32)              :: bc(2)
 
      type(sll_t_fft)           :: fw
      type(sll_t_fft)           :: bw
      real(f64), pointer     :: tmp(:)
      integer(i32), allocatable :: k_list(:)
      real(f64), allocatable :: z_r(:, :)
      real(f64), pointer     :: z_a(:, :)
      real(f64), allocatable :: mat(:, :)
      real(f64), allocatable :: cts(:)
      integer(i32), allocatable :: ipiv(:)
 
      real(f64) :: bc_coeffs_rmin(2:3) ! needed for Neumann at r=r_min
      real(f64) :: bc_coeffs_rmax(-2:-1) ! needed for Neumann at r=r_max
      integer   :: skip0                 ! needed for full circle
 
      type(sll_t_layout_2d), pointer :: layout_r
      type(sll_t_layout_2d), pointer :: layout_a
      type(sll_t_remap_plan_2d_real64), pointer :: rmp_ra
      type(sll_t_remap_plan_2d_real64), pointer :: rmp_ar
 
   end type sll_t_poisson_2d_polar_par
 
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
contains
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
   !=============================================================================
   subroutine sll_s_poisson_2d_polar_par_init(solver, &
                                              layout_r, &
                                              layout_a, &
                                              rmin, &
                                              rmax, &
                                              nr, &
                                              ntheta, &
                                              bc_rmin, &
                                              bc_rmax, &
                                              rgrid)
 
      type(sll_t_poisson_2d_polar_par), intent(out) :: solver
      type(sll_t_layout_2d), pointer       :: layout_r
      type(sll_t_layout_2d), pointer       :: layout_a
      real(f64), intent(in) :: rmin
      real(f64), intent(in) :: rmax
      integer(i32), intent(in) :: nr
      integer(i32), intent(in) :: ntheta
      integer(i32), intent(in) :: bc_rmin
      integer(i32), intent(in) :: bc_rmax
      real(f64), target, optional, intent(in) :: rgrid(:)
 
      character(len=*), parameter :: this_sub_name = 'sll_s_poisson_2d_polar_par_init'
 
      real(f64) :: hp, hm
      real(f64) :: inv_r
      real(f64) :: d1_coeffs(3)
      real(f64) :: d2_coeffs(3)
      real(f64), allocatable :: r_nodes(:)
 
      integer(i32)  :: i, j, k
      integer(i32)  :: bck(2)
      integer(i32)  :: last
      integer(i32)  :: sh
 
      integer(i32)  :: loc_sz_r(2) ! sequential in r direction
      integer(i32)  :: loc_sz_a(2) ! sequential in theta direction
      integer(i32)  :: glob_idx(2)
 
      integer(i32), allocatable :: k_list_glob(:)
 
      ! Consistency checks
      sll_assert_always(rmin >= 0.0_f64)
      sll_assert_always(rmin < rmax)
      sll_assert_always(nr >= 1)
      sll_assert_always(ntheta >= 1)
 
      if (bc_rmin == sll_p_polar_origin .and. rmin /= 0.0_f64) then
         sll_error(this_sub_name, "BC option 'sll_p_polar_origin' requires r_min = 0")
      end if
 
      ! Set boundary condition at r_min
      select case (bc_rmin)
      case (sll_p_dirichlet, sll_p_neumann, sll_p_neumann_mode_0, sll_p_polar_origin)
         solver%bc(1) = bc_rmin
      case default
         sll_error(this_sub_name, 'Unrecognized boundary condition at r_min')
      end select
 
      ! Set boundary condition at r_max
      select case (bc_rmax)
      case (sll_p_dirichlet, sll_p_neumann, sll_p_neumann_mode_0)
         solver%bc(2) = bc_rmax
      case default
         sll_error(this_sub_name, 'Unrecognized boundary condition at r_max')
      end select
 
      ! Important: if full circle is simulated, center point is not solved for!
      !            Therefore, solution is calculated on nr+1-skip0 points.
      sh = merge(1, 0, bc_rmin == sll_p_polar_origin)
 
      ! Consistency check: global size of 2D layouts must be (nr+1,ntheta)
      sll_assert_always(sll_o_get_layout_global_size_i(layout_r) == nr + 1 - sh)
      sll_assert_always(sll_o_get_layout_global_size_j(layout_r) == ntheta)
      !
      sll_assert_always(sll_o_get_layout_global_size_i(layout_a) == nr + 1 - sh)
      sll_assert_always(sll_o_get_layout_global_size_j(layout_a) == ntheta)
 
      ! Compute local size of 2D arrays in the two layouts
      call sll_o_compute_local_sizes(layout_r, loc_sz_r(1), loc_sz_r(2))
      call sll_o_compute_local_sizes(layout_a, loc_sz_a(1), loc_sz_a(2))
 
      ! Consistency check: layout_r sequential in r, layout_a sequential in theta
      sll_assert_always(loc_sz_r(1) == nr + 1 - sh)
      sll_assert_always(loc_sz_a(2) == ntheta)
 
      ! Store global information in solver
      solver%rmin = rmin
      solver%rmax = rmax
      solver%nr = nr
      solver%ntheta = ntheta
      solver%layout_a => layout_a
      solver%layout_r => layout_r
      solver%skip0 = sh
 
      ! r grid (possibly non-uniform)
      allocate (r_nodes(nr + 1))
 
      if (present(rgrid)) then  !--> Create computational grid from user data
 
         sll_assert_always(all(rgrid > 0.0_f64))
         if (bc_rmin == sll_p_polar_origin) then
            sll_assert_always(size(rgrid) == nr)
            sll_assert_always(rgrid(nr) == rmax)
            r_nodes(1) = -rgrid(1)
            r_nodes(2:) = rgrid(:)
         else
            sll_assert_always(size(rgrid) == nr + 1)
            sll_assert_always(rgrid(1) == rmin)
            sll_assert_always(rgrid(nr + 1) == rmax)
            r_nodes(:) = rgrid(:)
         end if
 
      else  !-------------------------> Create uniform grid
 
         if (bc_rmin == sll_p_polar_origin) then
            associate(rmin => rmax/real(2*nr + 1, f64))
               r_nodes(1) = -rmin
               call sll_s_new_array_linspace(r_nodes(2:), rmin, rmax, endpoint=.true.)
            end associate
         else
            call sll_s_new_array_linspace(r_nodes, rmin, rmax, endpoint=.true.)
         end if
 
      end if
 
      ! Allocate arrays global in r
      allocate (solver%z_r(nr + 1 + sh, loc_sz_r(2)))
      allocate (solver%mat((nr - 1)*3, loc_sz_r(2))) ! for each k, matrix depends on r
      allocate (solver%cts((nr - 1)*7))
      allocate (solver%ipiv(nr - 1))
 
      ! Remap objects between two layouts (for transposing between f_r and f_a)
      allocate (solver%z_a(loc_sz_a(1), ntheta))
      solver%rmp_ra => sll_o_new_remap_plan(solver%layout_r, solver%layout_a, solver%z_r)
      solver%rmp_ar => sll_o_new_remap_plan(solver%layout_a, solver%layout_r, solver%z_a)
      deallocate (solver%z_a)
 
      ! Allocate in ALIGNED fashion 1D work array for FFT
      solver%tmp => sll_f_fft_allocate_aligned_real(ntheta)
 
      ! Initialize plans for forward and backward FFTs
      call sll_s_fft_init_r2r_1d(solver%fw, &
                                 ntheta, &
                                 solver%tmp(:), &
                                 solver%tmp(:), &
                                 sll_p_fft_forward, &
                                 aligned=.true., &
                                 normalized=.true.)
 
      call sll_s_fft_init_r2r_1d(solver%bw, &
                                 ntheta, &
                                 solver%tmp(:), &
                                 solver%tmp(:), &
                                 sll_p_fft_backward, &
                                 aligned=.true., &
                                 normalized=.false.)
 
      ! Determine global k_list
      allocate (k_list_glob(ntheta))
      call sll_s_fft_get_k_list_r2r_1d(solver%fw, k_list_glob)
      ! Extract local k list
      allocate (solver%k_list(loc_sz_r(2)))
      do j = 1, loc_sz_r(2)
         glob_idx(:) = sll_o_local_to_global(solver%layout_r, [1, j])
         solver%k_list(j) = k_list_glob(glob_idx(2))
      end do
      deallocate (k_list_glob)
 
      ! Store matrix coefficients into solver%mat
      ! Cycle over k_j
      do j = 1, loc_sz_r(2)
 
         ! Get value of k_j from precomputed list of local values
         k = solver%k_list(j)
 
         !----------------------------------------------
         ! Compute boundary conditions type for mode k_j
         !----------------------------------------------
         bck(:) = solver%bc(:)
         do i = 1, 2
            if (bck(i) == sll_p_neumann_mode_0) then
               if (k == 0) then
                  bck(i) = sll_p_neumann
               else
                  bck(i) = sll_p_dirichlet
               end if
            end if
         end do
 
         !----------------------------------------------
         ! Compute matrix coefficients for a given k_j
         !----------------------------------------------
         do i = 2, nr
            hp = r_nodes(i + 1) - r_nodes(i)
            hm = r_nodes(i) - r_nodes(i - 1)
            inv_r = 1.0_f64/r_nodes(i)
            d1_coeffs(:) = [-hp/hm, (hp**2 - hm**2)/(hp*hm), hm/hp]/(hp + hm)
            d2_coeffs(:) = [2*hp/(hp + hm), -2.0_f64, 2*hm/(hp + hm)]/(hp*hm)
 
            solver%mat(3*(i - 1) - 2:3*(i - 1), j) = &
               -d2_coeffs(:) - d1_coeffs(:)*inv_r + [0.0_f64, (k*inv_r)**2, 0.0_f64]
         end do
 
         !----------------------------------------------
         ! Set boundary condition at rmin
         !----------------------------------------------
         if (bck(1) == sll_p_dirichlet) then ! Dirichlet
            solver%mat(1, j) = 0.0_f64
 
         else if (bck(1) == sll_p_neumann) then ! Neumann
 
            ! Coefficients of homogeneous boundary condition
            hp = r_nodes(3) - r_nodes(2)
            hm = r_nodes(2) - r_nodes(1)
            d1_coeffs(:) = [-2 - hp/hm, 2 + hp/hm + hm/hp, -hm/hp]
            solver%bc_coeffs_rmin(2:3) = -d1_coeffs(2:3)/d1_coeffs(1)
 
            ! Gaussian elimination: remove phi(1) variable using boundary condition
            solver%mat(3, j) = solver%mat(3, j) + solver%bc_coeffs_rmin(3)*solver%mat(1, j)
            solver%mat(2, j) = solver%mat(2, j) + solver%bc_coeffs_rmin(2)*solver%mat(1, j)
            solver%mat(1, j) = 0.0_f64
 
         else if (bck(1) == sll_p_polar_origin) then ! center of circular domain
 
            ! Gaussian elimination: phi(1) = (-1)^k phi(2)
            solver%mat(2, j) = solver%mat(2, j) + (-1)**k*solver%mat(1, j)
            solver%mat(1, j) = 0.0_f64
 
         end if
 
         !----------------------------------------------
         ! Set boundary condition at rmax
         !----------------------------------------------
         last = 3*(nr - 1)
         if (bck(2) == sll_p_dirichlet) then ! Dirichlet
            solver%mat(last, j) = 0.0_f64
 
         else if (bck(2) == sll_p_neumann) then ! Neumann
 
            ! Coefficients of homogeneous boundary condition
            hp = r_nodes(nr + 1) - r_nodes(nr)
            hm = r_nodes(nr) - r_nodes(nr - 1)
            d1_coeffs(:) = [hp/hm, -2 - hp/hm - hm/hp, 2 + hm/hp]
            solver%bc_coeffs_rmax(-2:-1) = -d1_coeffs(1:2)/d1_coeffs(3)
 
            ! Gaussian elimination: remove phi(last) variable using boundary condition
            solver%mat(last - 2, j) = solver%mat(last - 2, j) + solver%bc_coeffs_rmax(-2)*solver%mat(last, j)
            solver%mat(last - 1, j) = solver%mat(last - 1, j) + solver%bc_coeffs_rmax(-1)*solver%mat(last, j)
            solver%mat(last, j) = 0.0_f64
 
         end if
 
      end do
 
   end subroutine sll_s_poisson_2d_polar_par_init
 
   !=============================================================================
   subroutine sll_s_poisson_2d_polar_par_solve(solver, rho, phi)
      type(sll_t_poisson_2d_polar_par), intent(inout) :: solver
      real(f64), intent(in) :: rho(:, :)
      real(f64), target, intent(out) :: phi(:, :)
 
      integer(i32) :: nr, ntheta, bck(2)
      integer(i32) :: i, j, k
      integer(i32) :: nrpts
      integer(i32) :: sh
 
      nr = solver%nr
      ntheta = solver%ntheta
 
      ! Shift in radial grid indexing
      sh = solver%skip0 ! =1 if bc_rmin==sll_p_polar_origin, =0 otherwise
 
      ! Number of points in radial grid
      nrpts = nr + 1 - sh
 
      ! Consistency check: rho and phi must be given in layout sequential in theta
      call verify_argument_sizes_par(solver%layout_a, rho, 'rho')
      call verify_argument_sizes_par(solver%layout_a, phi, 'phi')
 
      ! Use output array 'phi' as 2D work array sequential in theta
      solver%z_a => phi(:, :)
 
      ! For each r_i, compute FFT of rho(r_i,theta) to obtain \hat{rho}(r_i,k)
      do i = 1, ubound(rho, 1)
         solver%tmp(:) = rho(i, :)
         call sll_s_fft_exec_r2r_1d(solver%fw, solver%tmp(:), solver%tmp(:))
         solver%z_a(i, :) = solver%tmp(:)
      end do
 
      ! Remap \hat{rho}(k) to layout distributed in k (global in r) -> \hat{rho}_k(r)
      call sll_o_apply_remap_2d(solver%rmp_ar, solver%z_a, solver%z_r)
 
      ! Cycle over k_j
      do j = 1, ubound(solver%z_r, 2)
 
         ! rhok(r) is k-th Fourier mode of rho(r,theta)
         ! phik(r) is k-th Fourier mode of phi(r,theta)
         ! rhok is 1D contiguous slice (column) of solver%z
         ! we will overwrite rhok with phik
         associate(rhok => solver%z_r(:, j), phik => solver%z_r(:, j))
 
            ! Solve tridiagonal system to obtain \hat{phi}_{k_j}(r) at internal points
            call sll_s_setup_cyclic_tridiag(solver%mat(:, j), nr - 1, solver%cts, solver%ipiv)
            call sll_o_solve_cyclic_tridiag(solver%cts, solver%ipiv, rhok(2 - sh:nrpts - 1), nr - 1, phik(2 - sh:nrpts - 1))
 
            ! Get value of k_j from precomputed list of local values
            k = solver%k_list(j)
 
            ! Compute boundary conditions type for mode k_j
            bck(:) = solver%bc(:)
            do i = 1, 2
               if (bck(i) == sll_p_neumann_mode_0) then
                  if (k == 0) then
                     bck(i) = sll_p_neumann
                  else
                     bck(i) = sll_p_dirichlet
                  end if
               end if
            end do
 
            ! Boundary condition at rmin
            if (bck(1) == sll_p_dirichlet) then ! Dirichlet
               phik(1) = 0.0_f64
            else if (bck(1) == sll_p_neumann) then ! Neumann
               associate(c => solver%bc_coeffs_rmin)
                  phik(1) = c(2)*phik(2) + c(3)*phik(3)
               end associate
            end if
 
            ! Boundary condition at rmax
            if (bck(2) == sll_p_dirichlet) then ! Dirichlet
               phik(nrpts) = 0.0_f64
            else if (bck(2) == sll_p_neumann) then ! Neumann
               associate(c => solver%bc_coeffs_rmax)
                  phik(nrpts) = c(-2)*phik(nrpts - 2) + c(-1)*phik(nrpts - 1)
               end associate
            end if
 
         end associate
 
      end do
 
      ! Redistribute \hat{phi}(r_i,k_j) into layout global in k
      call sll_o_apply_remap_2d(solver%rmp_ra, solver%z_r, solver%z_a)
 
      ! For each r_i, compute inverse FFT of \hat{phi}(r_i,k) to obtain phi(r_i,theta)
      do i = 1, ubound(solver%z_a, 1)
         solver%tmp(:) = solver%z_a(i, :)
         call sll_s_fft_exec_r2r_1d(solver%bw, solver%tmp(:), solver%tmp(:))
         phi(i, :) = solver%tmp(:)
      end do
 
   end subroutine sll_s_poisson_2d_polar_par_solve
 
   !=============================================================================
   subroutine sll_s_poisson_2d_polar_par_free(solver)
      type(sll_t_poisson_2d_polar_par), intent(inout) :: solver
 
      call sll_s_fft_free(solver%fw)
      call sll_s_fft_free(solver%bw)
 
      call sll_s_fft_deallocate_aligned_real(solver%tmp)
 
      deallocate (solver%z_r)
      deallocate (solver%mat)
      deallocate (solver%cts)
      deallocate (solver%ipiv)
 
      deallocate (solver%rmp_ra)
      deallocate (solver%rmp_ar)
 
      solver%layout_r => null()
      solver%layout_a => null()
      solver%rmp_ra => null()
      solver%rmp_ar => null()
 
   end subroutine sll_s_poisson_2d_polar_par_free
 
   !=============================================================================
   subroutine verify_argument_sizes_par(layout, array, array_name)
      type(sll_t_layout_2d), pointer    :: layout
      real(f64), intent(in) :: array(:, :)
      character(len=*), intent(in) :: array_name
 
      integer(i32)                :: n(2) ! nx_loc, ny_loc
      character(len=*), parameter :: sfmt = "('['(i0)','(i0)']')"
      character(len=256)          :: err_fmt
      character(len=256)          :: err_msg
 
      call sll_o_compute_local_sizes(layout, n(1), n(2))
 
      if (.not. all(shape(array) == n(:))) then
         ! Create error format string
         err_fmt = "(a,"//sfmt//",a,"//sfmt//",a)"
         ! Write error message to string
         write (err_msg, err_fmt) &
            "shape("//array_name//") = ", shape(array), &
            " is not compatible with local shape ", n, " of 2D layout"
         ! Stop execution with error
         sll_error("sll_s_poisson_2d_polar_par_solve", trim(err_msg))
      end if
 
   end subroutine verify_argument_sizes_par
 
end module sll_m_poisson_2d_polar_par