MODULE grid
! Grid module for spatial discretization
USE prec_const, ONLY: xp
USE basic, ONLY: allocate_array, speak
USE parallel, ONLY: my_id, comm_ky
USE iso_c_binding
IMPLICIT NONE
PRIVATE
! GRID parameters
INTEGER, PUBLIC, PROTECTED :: pmax = 2 ! The maximal Hermite-moment computed
INTEGER, PUBLIC, PROTECTED :: jmax = 1 ! The maximal Laguerre-moment computed
INTEGER, PUBLIC, PROTECTED :: dmax = 2 ! The maximal full GF set of i-moments v^dmax
INTEGER, PUBLIC, PROTECTED :: Nx = 64 ! Number of total internal grid points in x
REAL(xp), PUBLIC, PROTECTED :: Lx = 120_xp ! horizontal length of the spatial box
INTEGER, PUBLIC, PROTECTED :: Nexc = 1 ! factor to increase Lx when shear>0 (Lx = Nexc/kymin/shear)
INTEGER, PUBLIC, PROTECTED :: Ny = 64 ! Number of total internal grid points in y
REAL(xp), PUBLIC, PROTECTED :: Ly = 120_xp ! vertical length of the spatial box
INTEGER, PUBLIC, PROTECTED :: Nz = 1 ! Number of total perpendicular planes
! Modes number (determined from Nx and Ny)
INTEGER, PUBLIC, PROTECTED :: Nkx ! Number of total internal grid points in kx
INTEGER, PUBLIC, PROTECTED :: Nky ! Number of total internal grid points in ky
! Grid arrays
INTEGER, DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: parray, parray_full
INTEGER, DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: jarray, jarray_full
REAL(xp), DIMENSION(:,:), ALLOCATABLE, PUBLIC,PROTECTED :: kxarray ! moving kx grid (ExB)
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: kxarray0 ! inirial kyarray
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: kxarray_full
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: xarray
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: kyarray, kyarray_full
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: ikyarray, inv_ikyarray !mode indices arrays
REAL(xp), DIMENSION(:,:), ALLOCATABLE, PUBLIC,PROTECTED :: zarray
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: zarray_full
REAL(xp), DIMENSION(:,:,:,:), ALLOCATABLE, PUBLIC,PROTECTED :: kparray !moving kperp grid
REAL(xp), DIMENSION(:,:,:,:), ALLOCATABLE, PUBLIC,PROTECTED :: kparray0 !initial kperp grid
REAL(xp), DIMENSION(:,:,:,:), ALLOCATABLE, PUBLIC,PROTECTED :: kp2array !moving kperp^2 grid
! Kronecker delta for p=0, p=1, p=2, j=0, j=1
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: kroneck_p0, kroneck_p1, kroneck_p2, kroneck_p3
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: kroneck_j0, kroneck_j1
! Indexation variables
INTEGER, PUBLIC, PROTECTED :: ias ,iae ! species index
INTEGER, PUBLIC, PROTECTED :: ips ,ipe ! Hermite
INTEGER, PUBLIC, PROTECTED :: ijs ,ije ! Laguerre
INTEGER, PUBLIC, PROTECTED :: ikys,ikye ! Fourier y mode
INTEGER, PUBLIC, PROTECTED :: ikxs,ikxe ! Fourier x mode
INTEGER, PUBLIC, PROTECTED :: izs ,ize ! z-grid
INTEGER, PUBLIC, PROTECTED :: ieven, iodd ! indices for the staggered grids
INTEGER, PUBLIC, PROTECTED :: ip0, ip1, ip2, ip3, ij0, ij1, pp2
INTEGER, PUBLIC, PROTECTED :: ikx0, iky0, ikx_max, iky_max ! Indices of k-grid origin and max
! Total numbers of points for Hermite and Laguerre
INTEGER, PUBLIC, PROTECTED :: total_na
INTEGER, PUBLIC, PROTECTED :: total_np
INTEGER, PUBLIC, PROTECTED :: total_nj
INTEGER, PUBLIC, PROTECTED :: total_nky
INTEGER, PUBLIC, PROTECTED :: total_nkx
INTEGER, PUBLIC, PROTECTED :: total_nz
! Local numbers of points (without ghosts)
INTEGER, PUBLIC, PROTECTED :: local_na
INTEGER, PUBLIC, PROTECTED :: local_np
INTEGER, PUBLIC, PROTECTED :: local_nj
INTEGER, PUBLIC, PROTECTED :: local_nky
INTEGER, PUBLIC, PROTECTED :: local_nkx ! = total_Nkx = Nx, not parallel
INTEGER, PUBLIC, PROTECTED :: local_nx ! = local_nx_ptr from FFTW (used only for Fourier)
INTEGER, PUBLIC, PROTECTED :: local_nz
INTEGER, PUBLIC, PROTECTED :: local_nkp
INTEGER, PUBLIC, PROTECTED :: ngp, ngj, ngx, ngy, ngz ! number of ghosts points
INTEGER, PUBLIC, PROTECTED :: nzgrid ! one or two depending on the staggered grid option
! Local offsets
INTEGER, PUBLIC, PROTECTED :: local_na_offset
INTEGER, PUBLIC, PROTECTED :: local_np_offset
INTEGER, PUBLIC, PROTECTED :: local_nj_offset
INTEGER, PUBLIC, PROTECTED :: local_nky_offset
INTEGER, PUBLIC, PROTECTED :: local_nkx_offset
INTEGER, PUBLIC, PROTECTED :: local_nx_offset
INTEGER, PUBLIC, PROTECTED :: local_nz_offset
! C-pointer type for FFTW3
integer(C_INTPTR_T), PUBLIC,PROTECTED :: local_nx_ptr, local_nky_ptr
integer(C_INTPTR_T), PUBLIC,PROTECTED :: local_nx_ptr_offset, local_nky_ptr_offset
! Grid spacing and limits
INTEGER , PUBLIC, PROTECTED :: deltap
REAL(xp), PUBLIC, PROTECTED :: deltaz, inv_deltaz, inv_dkx
REAL(xp), PUBLIC, PROTECTED :: deltakx, deltaky, deltax, kx_max, ky_max, kx_min, ky_min!, kp_max
INTEGER , PUBLIC, PROTECTED :: local_pmin, local_pmax
INTEGER , PUBLIC, PROTECTED :: local_jmin, local_jmax
REAL(xp), PUBLIC, PROTECTED :: local_kymin, local_kymax
REAL(xp), PUBLIC, PROTECTED :: local_kxmin, local_kxmax
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC, PROTECTED :: local_zmin, local_zmax
! local z weights for computing simpson rule
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: zweights_SR
! Numerical diffusion scaling
REAL(xp), PUBLIC, PROTECTED :: diff_p_coeff, diff_j_coeff
REAL(xp), PUBLIC, PROTECTED :: diff_kx_coeff, diff_ky_coeff, diff_dz_coeff
LOGICAL, PUBLIC, PROTECTED :: SG = .true.! shifted grid flag
! Array to know the distribution of data among all processes (for MPI comm)
INTEGER, DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: counts_total_nkx, counts_nky, counts_nz
INTEGER, DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: displs_total_nkx, displs_nky, displs_nz
! Kperp array depends on kx, ky, z (geometry), eo (even or odd zgrid)
LOGICAL, PUBLIC, PROTECTED :: contains_kx0 = .false. ! flag if the proc contains kx=0 index
LOGICAL, PUBLIC, PROTECTED :: contains_ky0 = .false. ! flag if the proc contains ky=0 index
LOGICAL, PUBLIC, PROTECTED :: contains_kymax = .false. ! flag if the proc contains kx=kxmax index
LOGICAL, PUBLIC, PROTECTED :: contains_zmax = .false. ! flag if the proc contains z=pi-dz index
LOGICAL, PUBLIC, PROTECTED :: contains_zmin = .false. ! flag if the proc contains z=-pi index
LOGICAL, PUBLIC, PROTECTED :: SINGLE_KY = .false. ! to check if it is a single non 0 ky simulation
! Grid containing the polynomials degrees
LOGICAL, PUBLIC, PROTECTED :: CONTAINSp0, CONTAINSp1, CONTAINSp2, CONTAINSp3
LOGICAL, PUBLIC, PROTECTED :: SOLVE_POISSON, SOLVE_AMPERE
! Usefull inverse numbers
REAL(xp), PUBLIC, PROTECTED :: inv_Nx, inv_Ny, inv_Nz
! For Orszag filter
REAL(xp), PUBLIC, PROTECTED :: two_third_kxmax
REAL(xp), PUBLIC, PROTECTED :: two_third_kymax
REAL(xp), PUBLIC, PROTECTED :: two_third_kpmax
! 1D Antialiasing arrays (2/3 rule)
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: AA_x
REAL(xp), DIMENSION(:), ALLOCATABLE, PUBLIC,PROTECTED :: AA_y
! Public Functions
PUBLIC :: init_grids_data, set_grids, update_grids
PUBLIC :: set_kxgrid, set_kparray ! Can be called from geometry to include shear effects
PUBLIC :: grid_readinputs, grid_outputinputs
PUBLIC :: bar
! Precomputations
real(xp), PUBLIC, PROTECTED :: pmax_xp, jmax_xp
CONTAINS
SUBROUTINE grid_readinputs
! Read the input parameters
USE prec_const
IMPLICIT NONE
INTEGER :: lun = 90 ! File duplicated from STDIN
NAMELIST /GRID/ pmax, jmax, Nx, Lx, Ny, Ly, Nz, SG, Nexc
READ(lun,grid)
IF(Nz .EQ. 1) & ! overwrite SG option if Nz = 1 for safety of use
SG = .FALSE.
!! Compute the maximal degree of full GF moments set
! i.e. : all moments N_a^pj s.t. p+2j<=d are simulated (see GF closure)
dmax = min(pmax,2*jmax+1)
! Usefull precomputations
inv_Nx = 1._xp/REAL(Nx,xp)
inv_Ny = 1._xp/REAL(Ny,xp)
END SUBROUTINE grid_readinputs
!! Init the local and global number of points in all directions
SUBROUTINE init_grids_data(Na,EM,LINEARITY)
USE fourier, ONLY: init_grid_distr_and_plans
USE parallel, ONLY: num_procs_p, rank_p, num_procs_ky, rank_ky, num_procs_z, rank_z
USE utility, ONLY: decomp1D
IMPLICIT NONE
INTEGER, INTENT(IN) :: Na ! number of species coming from the model module
LOGICAL, INTENT(IN) :: EM ! electromagnetic effects (to skip odd Hermite or not)
CHARACTER(len=*), INTENT(IN) :: LINEARITY ! Linear or nonlinear run
INTEGER :: in, istart, iend
!!----------------- SPECIES INDICES (not parallelized)
ias = 1
iae = Na
total_Na = Na
local_Na = Na
local_Na_offset = ias - 1
!!----------------- HERMITE INDICES (parallelized)
! If no parallel dim (Nz=1) and no EM effects (beta=0), the moment hierarchy
!! is separable between odds and even P
IF((Nz .EQ. 1) .AND. .NOT. EM) THEN
deltap = 2
Ngp = 2 ! two ghosts cells for p+/-2 only
pp2 = 1 ! index p+2 is ip+1
ELSE
deltap = 1
Ngp = 4 ! four ghosts cells for p+/-1 and p+/-2 terms
pp2 = 2 ! index p+2 is ip+2
ENDIF
! Total number of Hermite polynomials we will evolve
total_np = (Pmax/deltap) + 1
! Local data distribution
CALL decomp1D(total_np, num_procs_p, rank_p, ips, ipe)
local_np = ipe - ips + 1
local_np_offset = ips - 1
! Allocate the grid arrays
ALLOCATE(parray_full(total_np))
ALLOCATE(parray(local_np+ngp))
!!----------------- LAGUERRE INDICES (not parallelized)
! Total number of J degrees
total_nj = jmax+1
local_jmax = jmax
Ngj= 2 ! 2-points ghosts for j+\-1 terms
! Indices of local data
ijs = 1; ije = jmax + 1
! Local number of J
local_nj = ije - ijs + 1
local_nj_offset = ijs - 1
! allocate global and local
ALLOCATE(jarray_full(total_nj))
ALLOCATE(jarray(local_nj+ngj))
ALLOCATE(kroneck_j0(local_nj+ngj));
ALLOCATE(kroneck_j1(local_nj+ngj));
!!----------------- SPATIAL GRIDS (parallelized)
!! Parallel distribution of kx ky grid
IF (LINEARITY .NE. 'linear') THEN ! we let FFTW distribute if we use it
IF (my_id .EQ. 0) write(*,*) 'FFTW3 y-grid distribution'
CALL init_grid_distr_and_plans(Nx,Ny,comm_ky,local_nx_ptr,local_nx_ptr_offset,local_nky_ptr,local_nky_ptr_offset)
ELSE ! otherwise we distribute equally
IF (my_id .EQ. 0) write(*,*) 'Manual y-grid distribution'
! balanced distribution among the processes
CALL decomp1D( Ny/2+1, num_procs_ky, rank_ky, ikys, ikye )
local_nky_ptr = ikye - ikys + 1
local_nky_ptr_offset = ikys - 1
ENDIF
!!----------------- BINORMAL KY INDICES (parallelized)
Nky = Ny/2+1 ! Defined only on positive kx since fields are real
total_nky = Nky
Ngy = 0 ! no ghosts cells in ky
ikys = INT(local_nky_ptr_offset + 1)
ikye = INT(ikys + local_nky_ptr - 1)
local_nky = ikye - ikys + 1
local_nky_offset = INT(local_nky_ptr_offset)
ALLOCATE(kyarray_full(Nky))
ALLOCATE(kyarray(local_nky))
ALLOCATE(ikyarray(Nky))
ALLOCATE(inv_ikyarray(Nky))
ALLOCATE(AA_y(local_nky))
!!---------------- RADIAL KX INDICES (not parallelized)
Nkx = Nx
total_nkx = Nx
ikxs = 1
ikxe = total_nkx
local_nkx = ikxe - ikxs + 1
local_nkx_offset = ikxs - 1
ALLOCATE(kxarray_full(total_nkx))
ALLOCATE(kxarray0(local_Nkx))
ALLOCATE(kxarray(local_nky,local_Nkx))
ALLOCATE(AA_x(local_nkx))
!!---------------- RADIAL X GRID (only for Fourier routines)
local_nx = INT(local_nx_ptr)
local_nx_offset = INT(local_nx_ptr_offset)
ALLOCATE(xarray(Nx))
!!---------------- PARALLEL Z GRID (parallelized)
total_nz = Nz
IF (SG) THEN
CALL speak('--2 staggered z grids--')
! indices for even p and odd p grids (used in kernel, jacobian, gij etc.)
ieven = 1
iodd = 2
nzgrid = 2
ELSE
ieven = 1
iodd = 1
nzgrid = 1
ENDIF
!! Parallel data distribution
IF( (Nz .EQ. 1) .AND. (num_procs_z .GT. 1) ) &
ERROR STOP '>> ERROR << Cannot have multiple core in z-direction (Nz = 1)'
! Local data distribution
CALL decomp1D(total_nz, num_procs_z, rank_z, izs, ize)
local_nz = ize - izs + 1
local_nz_offset = izs - 1
! Ghosts boundaries (depend on the order of z operators)
IF(Nz .EQ. 1) THEN
ngz = 0
ELSEIF(Nz .GE. 4) THEN
ngz =4
IF(mod(Nz,2) .NE. 0 ) THEN
ERROR STOP '>> ERROR << Nz must be an even number for Simpson integration rule !!!!'
ENDIF
ELSE
ERROR STOP '>> ERROR << Nz is not appropriate!!'
ENDIF
ALLOCATE(zarray(local_nz+ngz,nzgrid))
ALLOCATE(zarray_full(total_nz))
ALLOCATE(counts_nz (num_procs_z))
ALLOCATE(displs_nz (num_procs_z))
CALL allocate_array(local_zmax,1,nzgrid)
CALL allocate_array(local_zmin,1,nzgrid)
ALLOCATE(zweights_SR(local_nz))
! List of shift and local numbers between the different processes (used in scatterv and gatherv)
DO in = 0,num_procs_z-1
CALL decomp1D(total_nz, num_procs_z, in, istart, iend)
counts_nz(in+1) = iend-istart+1
displs_nz(in+1) = istart-1
ENDDO
!!---------------- Kperp grid
CALL allocate_array( kparray, 1,local_nky, 1,local_nkx, 1,local_nz+ngz, 1,nzgrid)
CALL allocate_array(kparray0, 1,local_nky, 1,local_nkx, 1,local_nz+ngz, 1,nzgrid)
CALL allocate_array(kp2array, 1,local_nky, 1,local_nkx, 1,local_nz+ngz, 1,nzgrid)
END SUBROUTINE
!!!! GRID REPRESENTATION
! We define the grids that contain ghosts (p,j,z) with indexing from 1 to Nlocal + nghost
! e.g. nghost = 4, nlocal = 4
! |x|x|_|_|_|_|x|x|
! 1 2 3 4 5 6 7 8 (x are ghosts)
! the interior points are contained between 1+Ng/2 and Nlocal+Ng/2
! The other grids are simply
! |_|_|_|_|
! 1 2 3 4
SUBROUTINE set_grids(shear,ExBrate,Npol,LINEARITY,N_HD,EM,Na)
REAL(xp), INTENT(IN) :: shear, ExBrate, Npol
CHARACTER(len=*), INTENT(IN) :: LINEARITY
INTEGER, INTENT(IN) :: N_HD
LOGICAL, INTENT(IN) :: EM
INTEGER, INTENT(IN) :: Na
CALL set_agrid(Na)
CALL set_pgrid(EM)
CALL set_jgrid
CALL set_kygrid(LINEARITY,N_HD)
CALL set_kxgrid(shear,ExBrate,Npol,1._xp,LINEARITY,N_HD) ! this will be redone after geometry if Cyq0_x0 .NE. 1
CALL set_xgrid
CALL set_zgrid (Npol)
END SUBROUTINE set_grids
SUBROUTINE set_agrid(Na) ! you're a sorcerer harry
IMPLICIT NONE
INTEGER, INTENT(IN) :: Na
ias = 1
iae = Na
total_Na = Na
local_Na = Na
local_Na_offset = ias - 1
END SUBROUTINE
SUBROUTINE set_pgrid(EM)
USE prec_const
IMPLICIT NONE
LOGICAL, INTENT(IN) :: EM
INTEGER :: ip
! Build the full grids on process 0 to diagnose it without comm
! P
DO ip = 1,total_np; parray_full(ip) = (ip-1)*deltap; END DO
!! local grid computations
! Fill pgrid array
DO ip = 1,local_np+ngp
parray(ip) = (ip-1-ngp/2+local_np_offset)*deltap
ENDDO
local_pmax = parray(local_np+ngp/2)
local_pmin = parray(1+ngp/2)
! Allocate kronecker arrays for p=0,1,2,3
ALLOCATE(kroneck_p0(local_np+ngp)); kroneck_p0 = 0._xp
ALLOCATE(kroneck_p1(local_np+ngp)); kroneck_p1 = 0._xp
ALLOCATE(kroneck_p2(local_np+ngp)); kroneck_p2 = 0._xp
ALLOCATE(kroneck_p3(local_np+ngp)); kroneck_p3 = 0._xp
DO ip = 1,local_np+ngp
SELECT CASE (parray(ip))
CASE(0)
ip0 = ip
kroneck_p0(ip) = 1._xp
CASE(1)
ip1 = ip
kroneck_p1(ip) = 1._xp
CASE(2)
ip2 = ip
kroneck_p2(ip) = 1._xp
CASE(3)
ip3 = ip
kroneck_p3(ip) = 1._xp
END SELECT
END DO
! Set local flags to avoid unnecessary logical operations
CONTAINSp0 = .FALSE.; CONTAINSp1 = .FALSE.
CONTAINSp2 = .FALSE.; CONTAINSp3 = .FALSE.
DO ip = 1,local_np+ngp
SELECT CASE (parray(ip))
CASE(0); CONTAINSp0 = .TRUE.
CASE(1); CONTAINSp1 = .TRUE.
CASE(2); CONTAINSp2 = .TRUE.
CASE(3); CONTAINSp3 = .TRUE.
END SELECT
END DO
! Flags that sets if Poisson and Ampere have to be solved locally
SOLVE_POISSON = .FALSE.; SOLVE_AMPERE = .FALSE.
DO ip = 1+ngp/2,local_np+ngp/2
SELECT CASE (parray(ip))
CASE(0); SOLVE_POISSON = .TRUE.
CASE(1); SOLVE_AMPERE = .TRUE.
END SELECT
END DO
! Precomputations
pmax_xp = real(pmax,xp)
diff_p_coeff = pmax_xp*(1._xp/(pmax_xp+1._xp))**6
! Overwrite SOLVE_AMPERE flag if beta is zero
IF(.NOT. EM) THEN
SOLVE_AMPERE = .FALSE.
ENDIF
END SUBROUTINE set_pgrid
SUBROUTINE set_jgrid
USE prec_const
IMPLICIT NONE
INTEGER :: ij
! Build the full grids on process 0 to diagnose it without comm
! J
DO ij = 1,total_nj
jarray_full(ij) = (ij-1)
END DO
DO ij = 1,local_nj+ngj
jarray(ij) = ij-1-ngj/2+local_nj_offset
END DO
local_jmax = jarray(local_nj+ngj/2)
local_jmin = jarray(1+ngj/2)
! Precomputations
jmax_xp = real(jmax,xp)
diff_j_coeff = jmax_xp*(1._xp/(jmax_xp+1._xp))**6
! j=0 and j=1 indices
DO ij = 1,local_nj+ngj
IF(jarray(ij) .EQ. 0) ij0 = ij
IF(jarray(ij) .EQ. 1) ij1 = ij
END DO
! Kronecker arrays for j
kroneck_j0 = 0._xp
kroneck_j1 = 0._xp
DO ij = 1,local_nj+ngj
SELECT CASE(jarray(ij))
CASE(0)
kroneck_j0(ij) = 1._xp
CASE(1)
kroneck_j1(ij) = 1._xp
END SELECT
END DO
END SUBROUTINE set_jgrid
SUBROUTINE set_kygrid(LINEARITY,N_HD)
USE prec_const
IMPLICIT NONE
CHARACTER(len=*), INTENT(IN) ::LINEARITY
INTEGER, INTENT(IN) :: N_HD
INTEGER :: iky
! Grid spacings
IF (Ny .EQ. 1) THEN
ERROR STOP "Gyacomo cannot run with only one ky"
ELSE
deltaky = 2._xp*PI/Ly
ky_max = (Nky-1)*deltaky
ky_min = deltaky
ENDIF
! Build the full grids on process 0 to diagnose it without comm
DO iky = 1,Nky
kyarray_full(iky) = REAL(iky-1,xp) * deltaky
! full indices grid (for ExB NL factor)
ikyarray(iky) = REAL((iky)-1,xp)
IF(ikyarray(iky) .GT. 0) THEN
inv_ikyarray(iky) = 1._xp/ikyarray(iky)
ELSE
inv_ikyarray(iky) = 0._xp
ENDIF
END DO
local_kymax = 0._xp
! Creating a grid ordered as dk*(0 1 2 3)
! We loop over the natural iky numbers (|1 2 3||4 5 6||... Nky|)
DO iky = 1,local_nky
! We shift the natural iky index by the offset to obtain the mpi dependent
! indexation (|1 2 3||1 2 3|... local_nky|)
IF(Ny .EQ. 1) THEN
kyarray(iky) = deltaky
kyarray_full(iky) = deltaky
SINGLE_KY = .TRUE.
ELSE
kyarray(iky) = kyarray_full(iky+local_nky_offset)
ENDIF
! Finding kx=0
IF (kyarray(iky) .EQ. 0) THEN
iky0 = iky
contains_ky0 = .true.
ENDIF
! Finding local kxmax value
IF (ABS(kyarray(iky)) .GT. local_kymax) THEN
local_kymax = ABS(kyarray(iky))
ENDIF
! Finding kxmax idx
IF (kyarray(iky) .EQ. ky_max) THEN
iky_max = iky
contains_kymax = .true.
ENDIF
END DO
! Orszag 2/3 filter
two_third_kymax = 2._xp/3._xp*deltaky*(Nky-1)
DO iky = 1,local_nky
IF ( (kyarray(iky) .LT. two_third_kymax) .OR. (LINEARITY .EQ. 'linear')) THEN
AA_y(iky) = 1._xp;
ELSE
AA_y(iky) = 0._xp;
ENDIF
END DO
! For hyperdiffusion
IF(LINEARITY.EQ.'linear') THEN
diff_ky_coeff= (1._xp/ky_max)**N_HD
ELSE
diff_ky_coeff= (1._xp/two_third_kymax)**N_HD
ENDIF
END SUBROUTINE set_kygrid
SUBROUTINE set_kxgrid(shear,ExBrate,Npol,Cyq0_x0,LINEARITY,N_HD)
USE prec_const, ONLY: xp, pi
USE basic, ONLY: speak, str
IMPLICIT NONE
REAL(xp), INTENT(IN) :: shear, Npol, Cyq0_x0, ExBrate
CHARACTER(len=*), INTENT(IN) ::LINEARITY
INTEGER, INTENT(IN) :: N_HD
INTEGER :: ikx, iky
REAL(xp):: Lx_adapted
IF(shear .GT. 0) THEN
CALL speak('Magnetic shear detected: set up sheared kx grid..')
! mininal size of box in x to respect dkx = 2pi shear dky
Lx_adapted = Ly/(2._xp*pi*shear*Npol*Cyq0_x0)
! Put Nexc to 0 so that it is computed from a target value Lx
IF(Nexc .EQ. 0) THEN
Nexc = CEILING(0.9 * Lx/Lx_adapted)
CALL speak('Adapted Nexc ='//str(Nexc))
ENDIF
! x length is adapted
Lx = Lx_adapted*Nexc
CALL speak('Simulation Lx ='//str(Lx))
ENDIF
deltakx = 2._xp*PI/Lx
inv_dkx = 1._xp/deltakx
IF(MODULO(total_nkx,2) .EQ. 0) THEN ! Even number of kx (-2 -1 0 1 2 3)
! Creating a grid ordered as dk*(0 1 2 3 -2 -1)
DO ikx = 1,total_nkx
kxarray_full(ikx) = deltakx*REAL(MODULO(ikx-1,total_nkx/2)-(total_nkx/2)*FLOOR(2.*real(ikx-1)/real(total_nkx)),xp)
IF (ikx .EQ. total_nkx/2+1) kxarray_full(ikx) = -kxarray_full(ikx)
END DO
kx_max = MAXVAL(kxarray_full)!(total_nkx/2)*deltakx
kx_min = MINVAL(kxarray_full)!-kx_max+deltakx
! Set local grid (not parallelized so same as full one)
local_kxmax = 0._xp
DO iky = 1,local_nky
DO ikx = 1,local_nkx
kxarray0(ikx) = kxarray_full(ikx+local_nkx_offset)
kxarray(iky,ikx) = kxarray_full(ikx+local_nkx_offset)
! Finding kx=0
IF (kxarray(1,ikx) .EQ. 0) THEN
ikx0 = ikx
contains_kx0 = .true.
ENDIF
! Finding local kxmax
IF (ABS(kxarray(iky,ikx)) .GT. local_kxmax) THEN
local_kxmax = ABS(kxarray(iky,ikx))
ikx_max = ikx
ENDIF
END DO
END DO
ELSE ! Odd number of kx (-2 -1 0 1 2)
ERROR STOP "Gyacomo is safer with an even Kx number"
ENDIF
! Orszag 2/3 filter
! We remove one point more if ExB is on since the moving grid
! can go up to kx=+-(kx_max+deltakx/2)
IF (ABS(ExBrate) .GT. 0) THEN
two_third_kxmax = 2._xp/3._xp*(kx_max)
ELSE
two_third_kxmax = 2._xp/3._xp*kx_max
ENDIF
! Antialiasing filter
DO iky = 1,local_nky
DO ikx = 1,local_nkx
IF ( ((kxarray(iky,ikx) .GT. -two_third_kxmax) .AND. &
(kxarray(iky,ikx) .LT. two_third_kxmax)) .OR. (LINEARITY .EQ. 'linear')) THEN
AA_x(ikx) = 1._xp;
ELSE
AA_x(ikx) = 0._xp;
ENDIF
END DO
END DO
! For hyperdiffusion
IF(LINEARITY.EQ.'linear') THEN
diff_kx_coeff= (1._xp/kx_max)**N_HD
ELSE
diff_kx_coeff= (1._xp/two_third_kxmax)**N_HD
ENDIF
END SUBROUTINE set_kxgrid
!----------- Radial x grid
! used only for the computation of the ExB shear NL factor
SUBROUTINE set_xgrid
USE prec_const, ONLY: xp, pi
IMPLICIT NONE
INTEGER :: ix
REAL(xp):: L_
L_ = 2._xp*pi/deltakx
deltax = L_/REAL(Nx,xp)
! full xgrid
DO ix = 1,Nx
xarray(ix) = REAL(ix-1,xp)*deltax
ENDDO
END SUBROUTINE set_xgrid
!----------- Parallel z grid
SUBROUTINE set_zgrid(Npol)
USE prec_const
IMPLICIT NONE
REAL(xp):: grid_shift, Lz, zmax, zmin, Npol
INTEGER :: iz, ig, eo
! Length of the flux tube (in ballooning angle)
Lz = 2._xp*pi*Npol
! Z stepping (#interval = #points since periodic)
deltaz = Lz/REAL(Nz,xp)
inv_deltaz = 1._xp/deltaz
! Parallel hyperdiffusion coefficient
diff_dz_coeff = (deltaz/2._xp)**4 ! adaptive fourth derivative (~GENE)
IF (SG) THEN
grid_shift = deltaz/2._xp
ELSE
grid_shift = 0._xp
ENDIF
! Build the full grids on process 0 to diagnose it without comm
IF (Nz .EQ. 1) &
Npol = 0._xp
zmax = 0; zmin = 0;
DO iz = 1,total_nz ! z in [-pi pi-dz] x Npol
zarray_full(iz) = REAL(iz-1,xp)*deltaz - Lz/2._xp
IF(zarray_full(iz) .GT. zmax) zmax = zarray_full(iz)
IF(zarray_full(iz) .LT. zmin) zmin = zarray_full(iz)
END DO
IF(Nz .EQ. 1) &
zarray_full(izs) = 0
! Local z array
DO iz = 1,local_nz
DO eo = 1,nzgrid
zarray(iz+ngz/2,eo) = zarray_full(iz+local_nz_offset) + REAL(eo-1,xp)*grid_shift
ENDDO
ENDDO
DO eo = 1,nzgrid
! Find local extrema
local_zmax(eo) = zarray(local_nz+ngz/2,eo)
local_zmin(eo) = zarray(1+ngz/2,eo)
! continue in z<pi and z>pi
DO ig = 1,ngz/2
zarray(local_nz+ngz/2+ig,eo) = zarray(local_nz+ngz/2,eo) + ig*deltaz
zarray(ig,eo) = zarray(1+ngz/2,eo) - (3-ig)*deltaz
ENDDO
! Set up the flags to know if the process contains the tip and/or the tail
! of the z domain (important for z-boundary condition)
IF(abs(local_zmin(eo) - (zmin+REAL(eo-1,xp)*grid_shift)) .LT. EPSILON(zmin)) &
contains_zmin = .TRUE.
IF(abs(local_zmax(eo) - (zmax+REAL(eo-1,xp)*grid_shift)) .LT. EPSILON(zmax)) &
contains_zmax = .TRUE.
ENDDO
! local weights for Simpson rule
IF(total_nz .EQ. 1) THEN
zweights_SR = 1._xp
ELSE
DO iz = 1,local_nz
IF(MODULO(iz+local_nz_offset,2) .EQ. 1) THEN ! odd iz
zweights_SR(iz) = onethird*deltaz*2._xp
ELSE ! even iz
zweights_SR(iz) = onethird*deltaz*4._xp
ENDIF
ENDDO
ENDIF
END SUBROUTINE set_zgrid
SUBROUTINE set_kparray(gxx, gxy, gyy, inv_hatB2)
IMPLICIT NONE
REAL(xp), DIMENSION(local_nz+ngz,nzgrid), INTENT(IN) :: gxx,gxy,gyy,inv_hatB2
INTEGER :: eo,iz,iky,ikx
REAL(xp) :: kx, ky
DO eo = 1,nzgrid
DO iz = 1,local_nz+ngz
DO iky = 1,local_nky
ky = kyarray(iky)
DO ikx = 1,local_nkx
kx = kxarray(iky,ikx)
! there is a factor 1/B from the normalization; important to match GENE
! this factor comes from $b_a$ argument in the Bessel. Kperp is not used otherwise.
kp2array(iky, ikx, iz, eo) = &
(gxx(iz,eo)*kx**2 + 2._xp*gxy(iz,eo)*kx*ky + gyy(iz,eo)*ky**2)*inv_hatB2(iz,eo)
kparray (iky, ikx, iz, eo) = SQRT(kp2array(iky, ikx, iz, eo))
kparray0(iky, ikx, iz, eo) = SQRT(kp2array(iky, ikx, iz, eo))
ENDDO
ENDDO
ENDDO
ENDDO
two_third_kpmax = 2._xp/3._xp * MAXVAL(kparray)
END SUBROUTINE
SUBROUTINE update_grids (sky_ExB,gxx,gxy,gyy,inv_hatB2)
IMPLICIT NONE
REAL(xp), DIMENSION(local_nky),INTENT(IN) :: sky_ExB ! ExB correction dkx = gamma_E ky dtshift
REAL(xp), DIMENSION(local_nz+ngz,nzgrid), INTENT(IN) :: gxx,gxy,gyy,inv_hatB2
INTEGER :: eo,iz,iky,ikx
REAL(xp) :: kx, ky
! Update the kx grid
DO ikx = 1,total_Nkx
DO iky = 1,local_nky
! For positive ExBrate, sky_ExB is negative
kxarray(iky,ikx) = kxarray0(ikx) + sky_ExB(iky)
ENDDO
ENDDO
! Update the kperp grid
DO eo = 1,nzgrid
DO iz = 1,local_nz+ngz
DO ikx = 1,local_nkx
DO iky = 1,local_nky
ky = kyarray(iky)
kx = kxarray(iky,ikx)
kp2array(iky,ikx,iz,eo) = (gxx(iz,eo)*kx**2 + 2._xp*gxy(iz,eo)*kx*ky + gyy(iz,eo)*ky**2)*inv_hatB2(iz,eo)
ENDDO
ENDDO
ENDDO
ENDDO
END SUBROUTINE
SUBROUTINE grid_outputinputs(fid)
! Write the input parameters to the results_xx.h5 file
USE futils, ONLY: attach, creatd
IMPLICIT NONE
INTEGER, INTENT(in) :: fid
CHARACTER(len=256) :: str
WRITE(str,'(a)') '/data/input/grid'
CALL creatd(fid, 0,(/0/),TRIM(str),'Grid Input')
CALL attach(fid, TRIM(str), "pmax", pmax)
CALL attach(fid, TRIM(str),"deltap", deltap)
CALL attach(fid, TRIM(str), "Np", total_np)
CALL attach(fid, TRIM(str), "jmax", jmax)
CALL attach(fid, TRIM(str), "Nj", total_nj)
CALL attach(fid, TRIM(str), "Nkx", Nkx)
CALL attach(fid, TRIM(str), "Nx", Nx)
CALL attach(fid, TRIM(str), "Lx", Lx)
CALL attach(fid, TRIM(str), "Nexc", Nexc)
CALL attach(fid, TRIM(str), "Ny", Ny)
CALL attach(fid, TRIM(str), "Nky", Nky)
CALL attach(fid, TRIM(str), "Ly", Ly)
CALL attach(fid, TRIM(str), "Nz", Nz)
CALL attach(fid, TRIM(str), "total_nkx", total_nkx)
CALL attach(fid, TRIM(str), "Nky", Nky)
CALL attach(fid, TRIM(str), "SG", SG)
END SUBROUTINE grid_outputinputs
FUNCTION bar(p_,j_)
IMPLICIT NONE
INTEGER :: bar, p_, j_
bar = (jmax+1)*p_ + j_ + 1
END FUNCTION
END MODULE grid