!> Set up the perpendicular wavenumbers by calling the appropriate sub-modules. module kt_grids use abstract_config, only: abstract_config_type, CONFIG_MAX_NAME_LEN implicit none private public :: init_kt_grids, is_box, finish_kt_grids, check_kt_grids, wnml_kt public :: aky, theta0, akx, naky, ntheta0, nx, ny, enforce_reality, ikx, jtwist public :: kwork_filter, kperp2, inv_kperp2 public :: itleft, itright, l_links, r_links, n_links, get_leftmost_it, get_rightmost_it public :: supercell_labels, n_supercells, n_cells, flux_tube_field_like_to_ballooning_space public :: kt_grids_config_type, set_kt_grids_config, get_kt_grids_config logical, dimension(:,:), allocatable :: kwork_filter integer, dimension(:,:), allocatable :: itleft, itright, l_links, r_links integer, dimension(:,:,:), allocatable :: n_links !> Labels the supercells as a function of {it, ik}. Primarily !> to help with post-processing, although can also help with !> conversion from flux tube to ballooning space. This probably !> really belongs in kt_grids, but placing here alongside related !> itleft, itright etc. integer, dimension(:, :), allocatable :: supercell_labels !> Records the number of unique supercells for each ky, integer, dimension(:), allocatable :: n_supercells !> Records the number of cells which are members of the supercell !> which owns the particular {it, ik} point. One would expect this !> to be equivalent to l_links + r_links + 1. integer, dimension(:, :), allocatable :: n_cells real, dimension (:,:,:), allocatable :: kperp2, inv_kperp2 !> The \(\theta_0(k_x, k_y)\) grid real, dimension (:,:), allocatable :: theta0 !> The \(k_x \rho\) grid real, dimension (:), allocatable :: akx !> The \(k_y \rho\) grid real, dimension (:), allocatable :: aky !> Discrete kx grid index integer, dimension(:), allocatable :: ikx !> Number of \(k_y \rho\) points integer :: naky !> Number of \(\theta_0\) points integer :: ntheta0 !> Number of (real space) \(x\) points integer :: nx !> Number of (real space) \(y\) points integer :: ny !> Related to the box size in x. See !> [[kt_grids_box_parameters:jtwist]] integer :: jtwist !> The type of perpendicular wavenumber grid used. See !> [[kt_grids_knobs:grid_option]] character(20) :: grid_option !> The rhostar (`rhoref/Lref`) to use to calculate ky from `n0`. Only used if `n0` also set !> greater than zero. real :: rhostar ! internal variables integer :: gridopt_switch integer, parameter :: gridopt_single = 1, gridopt_range = 2, & gridopt_specified = 3, gridopt_box = 4 logical :: enforce_reality = .false., is_box = .false., initialized = .false., kp2init=.false. !> Used to represent the input configuration of kt_grids type, extends(abstract_config_type) :: kt_grids_config_type ! namelist : kt_grids_knobs ! indexed : false !> Controls the type of perpendicular wavenumber grid to use. !> Can be one of !> !> - 'single' Evolve a single Fourier mode. Set up !> with [[kt_grids_single_parameters]]. !> - 'default' Same as 'single' !> - 'range' Evolve a range of equally spaced Fourier !> modes. Set up with [[kt_grids_range_parameters]]. !> - 'specified' Evolve an arbitrary set of Fourier modes. Set !> up with [[kt_grids_specified_parameters]] and corresponding !> [[kt_grids_specified_element]]. !> - 'box' Simulation domain is logically rectangular. Set up with !> [[kt_grids_box_parameters]]. Required for nonlinear simulations. !> - 'nonlinear' Same as 'box.' !> - 'xbox' Experimental. !> character(len = 20) :: grid_option = "default" !> The rhostar (`rhoref/Lref`) to use to calculate ky from `n0`. !> Only used if `n0` also set in the active kt_grids sub-type real :: rhostar = 1.0e-4 contains procedure, public :: read => read_kt_grids_config procedure, public :: write => write_kt_grids_config procedure, public :: reset => reset_kt_grids_config procedure, public :: broadcast => broadcast_kt_grids_config procedure, public, nopass :: get_default_name => get_default_name_kt_grids_config procedure, public, nopass :: get_default_requires_index => get_default_requires_index_kt_grids_config end type kt_grids_config_type type(kt_grids_config_type) :: kt_grids_config contains !> FIXME : Add documentation subroutine init_kt_grids(kt_grids_config_in, kt_grids_single_config_in, & kt_grids_range_config_in, kt_grids_specified_config_in, & kt_grids_specified_element_config_in, kt_grids_box_config_in) use kt_grids_single, only: init_kt_grids_single, get_sizes_and_grids_single use kt_grids_single, only: kt_grids_single_config_type use kt_grids_range, only: init_kt_grids_range, get_sizes_and_grids_range use kt_grids_range, only: kt_grids_range_config_type use kt_grids_specified, only: init_kt_grids_specified, get_sizes_and_grids_specified use kt_grids_specified, only: kt_grids_specified_config_type, kt_grids_specified_element_config_type use kt_grids_box, only: init_kt_grids_box, get_sizes_and_grids_box, jtwist_box => jtwist use kt_grids_box, only: kt_grids_box_config_type implicit none type(kt_grids_config_type), intent(in), optional :: kt_grids_config_in type(kt_grids_single_config_type), intent(in), optional :: kt_grids_single_config_in type(kt_grids_range_config_type), intent(in), optional :: kt_grids_range_config_in type(kt_grids_specified_config_type), intent(in), optional :: kt_grids_specified_config_in type(kt_grids_specified_element_config_type), intent(in), dimension(:), optional :: kt_grids_specified_element_config_in type(kt_grids_box_config_type), intent(in), optional :: kt_grids_box_config_in if (initialized) return initialized = .true. call read_parameters(kt_grids_config_in) select case (gridopt_switch) case (gridopt_single) call init_kt_grids_single(rhostar, kt_grids_single_config_in) call get_sizes_and_grids_single(aky, theta0, akx, ikx, naky, ntheta0, nx, ny) case (gridopt_range) call init_kt_grids_range(rhostar, kt_grids_range_config_in) call get_sizes_and_grids_range(aky, theta0, akx, ikx, naky, ntheta0, nx, ny) case (gridopt_specified) ! Note specified doesn't support setting from n0 so doesn't take rhostar call init_kt_grids_specified(kt_grids_specified_config_in) call get_sizes_and_grids_specified(aky, theta0, akx, ikx, naky, ntheta0, nx, ny, & kt_grids_specified_element_config_in) case (gridopt_box) call init_kt_grids_box(rhostar, kt_grids_box_config_in) call get_sizes_and_grids_box(aky, theta0, akx, ikx, naky, ntheta0, nx, ny) enforce_reality = .true. is_box = .true. end select jtwist = jtwist_box allocate(kwork_filter(ntheta0,naky)) kwork_filter=.false. call init_kperp2 allocate (itleft(ntheta0, naky), itright(ntheta0, naky)) call compute_itleft_and_itright(compute_jshift0(), itleft, itright) allocate (l_links(ntheta0, naky), r_links(ntheta0, naky), n_links(2, ntheta0, naky)) call count_links(itleft, itright, l_links, r_links, n_links) allocate(supercell_labels(ntheta0, naky), n_supercells(naky), n_cells(ntheta0, naky)) call calculate_supercell_labels(supercell_labels, n_supercells, n_cells) end subroutine init_kt_grids !> FIXME : Add documentation subroutine read_parameters(kt_grids_config_in) use file_utils, only: error_unit use text_options, only: text_option, get_option_value implicit none type(kt_grids_config_type), intent(in), optional :: kt_grids_config_in type (text_option), dimension (6), parameter :: gridopts = & [ text_option('default', gridopt_single), & text_option('single', gridopt_single), & text_option('range', gridopt_range), & text_option('specified', gridopt_specified), & text_option('box', gridopt_box), & text_option('nonlinear', gridopt_box) ] integer :: ierr if (present(kt_grids_config_in)) kt_grids_config = kt_grids_config_in call kt_grids_config%init(name = 'kt_grids_knobs', requires_index = .false.) ! Copy out internal values into module level parameters associate(self => kt_grids_config) #include "kt_grids_copy_out_auto_gen.inc" end associate ierr = error_unit() call get_option_value (grid_option, gridopts, gridopt_switch, & ierr, "grid_option in kt_grids_knobs",.true.) end subroutine read_parameters !> FIXME : Add documentation subroutine wnml_kt(unit) use kt_grids_single, only: wnml_kt_grids_single use kt_grids_range, only: wnml_kt_grids_range use kt_grids_specified, only: wnml_kt_grids_specified use kt_grids_box, only: wnml_kt_grids_box implicit none integer, intent(in) :: unit call kt_grids_config%write(unit) select case (gridopt_switch) case (gridopt_single) call wnml_kt_grids_single (unit) case (gridopt_range) call wnml_kt_grids_range (unit) case (gridopt_specified) call wnml_kt_grids_specified (unit) case (gridopt_box) call wnml_kt_grids_box (unit) end select end subroutine wnml_kt !> FIXME : Add documentation subroutine init_kperp2 use theta_grid, only: ntgrid, gds2, gds21, gds22, shat use warning_helpers, only: is_zero implicit none integer :: ik, it, ig if (kp2init) return kp2init = .true. allocate (kperp2(-ntgrid:ntgrid,ntheta0,naky)) do ik = 1, naky if (is_zero(shat)) then do it = 1, ntheta0 kperp2(:,it,ik) = akx(it)*akx(it) + aky(ik)*aky(ik) end do else if (is_zero(aky(ik))) then do it = 1, ntheta0 kperp2(:,it,ik) = akx(it)*akx(it)*gds22/(shat*shat) end do else do it = 1, ntheta0 kperp2(:,it,ik) = aky(ik)*aky(ik) & *(gds2 + 2.0*theta0(it,ik)*gds21 & + theta0(it,ik)*theta0(it,ik)*gds22) end do end if end if end do allocate (inv_kperp2(-ntgrid:ntgrid,ntheta0,naky)) do ik = 1, naky do it = 1, ntheta0 do ig = -ntgrid, ntgrid if (abs(kperp2(ig, it, ik)) > epsilon(0.0)) then inv_kperp2(ig, it, ik) = 1.0/kperp2(ig, it, ik) else inv_kperp2(ig, it, ik) = 0.0 end if end do end do end do end subroutine init_kperp2 !> FIXME : Add documentation subroutine finish_kt_grids use kt_grids_single, only: finish_kt_grids_single use kt_grids_range, only: finish_kt_grids_range use kt_grids_specified, only: finish_kt_grids_specified use kt_grids_box, only: finish_kt_grids_box implicit none if (allocated(aky)) deallocate (akx, aky, theta0, ikx) if (allocated(kwork_filter)) deallocate(kwork_filter) if (allocated(kperp2)) deallocate(kperp2) if (allocated(inv_kperp2)) deallocate(inv_kperp2) if (allocated(itleft)) deallocate(itleft, itright) if (allocated(l_links)) deallocate(l_links, r_links, n_links) if (allocated(supercell_labels)) deallocate (supercell_labels, n_cells, n_supercells) enforce_reality = .false. ; is_box = .false. kp2init = .false. ; initialized = .false. call finish_kt_grids_single call finish_kt_grids_specified call finish_kt_grids_range call finish_kt_grids_box call kt_grids_config%reset() end subroutine finish_kt_grids !> Compute the spacing in kx indices between connected kx domains !> (cells) for the lowest non-zonal ky. integer function compute_jshift0() result(jshift0) use theta_grid, only: ntgrid, theta use theta_grid, only: shat use mp, only: mp_abort use warning_helpers, only: is_not_zero, is_zero implicit none real :: theta_extent, theta0_spacing ! jshift0 corresponds to J (not delta j) from Beer thesis (unsure about +0.01) -- MAB ! delta j is the number of akxs (i.e. theta0s) 'skipped' when connecting one ! parallel segment to the next to satisfy the twist and shift ! boundary conditions. delta j = (ik - 1) * jshift0 . EGH if (is_zero(shat)) then jshift0 = 0 return end if theta_extent = theta(ntgrid) - theta(-ntgrid) if (naky > 1 .and. ntheta0 > 1) then theta0_spacing = theta0(2, 2) - theta0(1, 2) else if (naky == 1 .and. ntheta0 > 1 .and. is_not_zero(aky(1))) then ! Not clear we _ever_ get here as box mode usually requires ! aky(1) == 0. One could imagine potentially using linked boundary ! conditions and kt_grids_range to end up here. theta0_spacing = theta0(2, 1) - theta0(1, 1) else ! Not possible to work out a theta0 spacing here ! so just set up such that jshift0 == 1 theta0_spacing = theta_extent end if ! Work out how many theta0 "grid points" we must step over ! to travel a distance equiavlent from one end of a cell ! to the other end. jshift0 = nint(theta_extent/theta0_spacing) ! Sanity check the setup. if (abs(jshift0 * theta0_spacing - theta_extent) > 100*epsilon(0.0)) then call mp_abort('Calculation of jshift0 indicates inconsistent grids.', .true.) end if end function compute_jshift0 !> For the passed jshift0 value determine the it (kx) indices to !> the left and right of each {it, ik} index pair. Used to construct !> the connected boundary conditions. subroutine compute_itleft_and_itright(jshift0, itleft, itright) use warning_helpers, only: is_not_zero implicit none integer, intent(in) :: jshift0 integer, dimension(:, :), intent(out) :: itleft, itright integer :: ik, it, it0, itl, itr ! Note the below appears to assume ik = 1 is always the ky = 0 ! mode but this is not consistent with potentially more general ! treatement in the init_connected_bc related code. We also revisit ! ik = 1 in the below loop so this initialisation is likely ! overwritten in all cases. itleft(:, 1) = -1 ! No connected segments when ky = 0. EGH itright(:, 1) = -1 !$OMP PARALLEL DO DEFAULT(none) & !$OMP PRIVATE(ik, it, it0, itl, itr) & !$OMP SHARED(naky, ntheta0, aky, jshift0, itleft, itright) & !$OMP COLLAPSE(2) & !$OMP SCHEDULE(static) do ik = 1, naky do it = 1, ntheta0 if (it > (ntheta0+1)/2) then ! akx is negative (-1 shift because akx(1) = 0) -- MAB it0 = it - ntheta0 - 1 else ! akx is positive (-1 shift because akx(1) = 0) -- MAB it0 = it - 1 end if if (ik == 1) then if (is_not_zero(aky(ik)) .and. naky == 1) then ! for this case, jshift0 is delta j from Beer thesis -- MAB itl = it0 + jshift0 itr = it0 - jshift0 else itl = it0 itr = it0 end if else ! ik = 1 corresponds to aky=0, so shift index by 1 ! (ik-1)*jshift0 is delta j from Beer thesis -- MAB itl = it0 + (ik-1)*jshift0 itr = it0 - (ik-1)*jshift0 end if ! remap to usual GS2 indices -- MAB if (itl >= 0 .and. itl < (ntheta0+1)/2) then itleft(it, ik) = itl + 1 else if (itl + ntheta0 + 1 > (ntheta0+1)/2 & .and. itl + ntheta0 + 1 <= ntheta0) then itleft(it, ik) = itl + ntheta0 + 1 else ! for periodic, need to change below -- MAB/BD ! j' = j + delta j not included in simulation, so can't connect -- MAB itleft(it, ik) = -1 end if ! same stuff for j' = j - delta j -- MAB if (itr >= 0 .and. itr < (ntheta0+1)/2) then itright(it, ik) = itr + 1 else if (itr + ntheta0 + 1 > (ntheta0+1)/2 & .and. itr + ntheta0 + 1 <= ntheta0) then itright(it, ik) = itr + ntheta0 + 1 else ! for periodic, need to change below -- MAB/BD itright(it, ik) = -1 end if end do end do !$OMP END PARALLEL DO end subroutine compute_itleft_and_itright !> Count how many links are to the left/right of each {it,ik} point subroutine count_links_one_side(it_connections, links) use mp, only: mp_abort implicit none integer, dimension(:, :), intent(in) :: it_connections !> Links is the number of links to the left or right integer, dimension(:, :), intent(out) :: links integer :: it, ik, it_star, it_star_next !> Arbitrary limit on the maximum number of allowed links in connected boundary conditions integer, parameter :: max_allowed_links = 5000 !$OMP PARALLEL DO DEFAULT(none) & !$OMP PRIVATE(ik, it, it_star, it_star_next) & !$OMP SHARED(naky, ntheta0, links, it_connections) & !$OMP COLLAPSE(2) & !$OMP SCHEDULE(static) do ik = 1, naky do it = 1, ntheta0 links(it, ik) = 0 it_star = it loop_over_connections: do it_star_next = it_connections(it_star, ik) ! If connected to self (periodic) exit if (it_star == it_star_next) exit loop_over_connections ! If no more connections exit if (it_star_next < 0) exit loop_over_connections links(it, ik) = links(it, ik) + 1 it_star = it_star_next if (links(it, ik) > max_allowed_links) then call mp_abort('links error',.true.) endif end do loop_over_connections end do end do !$OMP END PARALLEL DO end subroutine count_links_one_side !> Count the links to the left and right of each {it, ik} point !> and work out how many values are required to set the bc. subroutine count_links(itleft, itright, l_links, r_links, n_links) implicit none integer, dimension(:, :), intent(in) :: itleft, itright integer, dimension(:, :), intent(out) :: l_links, r_links integer, dimension(:, :, :), intent(out) :: n_links ! First count the number of links to the left and right call count_links_one_side(itleft, l_links) call count_links_one_side(itright, r_links) ! Then work out the number of values needed for the bc. ! 'n_links' complex numbers are needed to specify bc for (it, ik) region ! ignoring wfb ! n_links(1,:,:) is for v_par > 0, etc. where (l_links == 0) n_links(1, :, :) = 0 elsewhere n_links(1, :, :) = 2 * l_links - 1 end where where (r_links == 0) n_links(2, :, :) = 0 elsewhere n_links(2, :, :) = 2 * r_links - 1 end where end subroutine count_links !> Helper function for finding the leftmost it of supercell elemental integer function get_leftmost_it(it,ik) result(it_cur) implicit none integer, intent(in) :: it, ik it_cur = it do while (itleft(it_cur, ik) >= 0 .and. itleft(it_cur, ik) /= it_cur) !Keep updating it_cur with left connected it until there are no !connections to left it_cur = itleft(it_cur, ik) end do end function get_leftmost_it !> Helper function for finding the rightmost it of supercell elemental integer function get_rightmost_it(it, ik) result(it_cur) implicit none integer, intent(in) :: it, ik !If not linked then no connections so only one cell in supercell it_cur = it do while (itright(it_cur, ik) >= 0 .and. itright(it_cur, ik) /= it_cur) !Keep updating it_cur with right connected it until there are no !connections to right it_cur = itright(it_cur, ik) end do end function get_rightmost_it !> Assigns a label to each {it, ik} point denoting which unique !> supercell the point belongs to. subroutine calculate_supercell_labels(supercell_labels, n_supercells, n_cells) implicit none integer, dimension(:, :), intent(out) :: supercell_labels, n_cells integer, dimension(:), intent(out) :: n_supercells integer, dimension(:), allocatable :: temporary_it_leftmost integer :: it, ik, isc, ncell supercell_labels = -1 n_supercells = 0 n_cells = 0 isc = -1 allocate(temporary_it_leftmost(ntheta0)) temporary_it_leftmost = -1 do ik = 1, naky do it = 1, ntheta0 temporary_it_leftmost(it) = get_leftmost_it(it, ik) end do do it = 1, ntheta0 if (temporary_it_leftmost(it) >= 0) then ! Note ncell isn't currently used but records the number ! of cells for the current supercell. We struggle to ! store it efficiently at this point as we don't yet know ! how many supercells there are. We currently store it in ! an array of shape [ntheta0, naky] so that we can look ! up the number of cells given {it, ik} rather than ! {isc}, so there is some duplication. ncell = count(temporary_it_leftmost == temporary_it_leftmost(it)) isc = isc + 1 where (temporary_it_leftmost == temporary_it_leftmost(it)) temporary_it_leftmost = -1 supercell_labels(:, ik) = isc n_cells(:, ik) = ncell end where n_supercells(ik) = n_supercells(ik) + 1 end if end do end do end subroutine calculate_supercell_labels !> Extracts the field like data corresponding to the supercell with !> the {it, ik} wavenumber pair as a member and constructs the !> ballooning space form. By default includes duplicate theta points !> but can drop these on request. subroutine flux_tube_field_like_to_ballooning_space(it, ik, field_like_ft, & theta_bs, field_like_bs, drop_duplicates) use theta_grid, only: ntgrid, theta use sorting, only: quicksort use optionals, only: get_option_with_default implicit none integer, intent(in) :: it, ik complex, dimension(-ntgrid:, :, :), intent(in) :: field_like_ft real, dimension(:), allocatable, intent(out) :: theta_bs complex, dimension(:), allocatable, intent(out) :: field_like_bs logical, intent(in), optional :: drop_duplicates real, dimension(:), allocatable :: theta0_values real, dimension(:), allocatable :: theta0_indices_real integer, dimension(:), allocatable :: theta0_indices logical, dimension(:), allocatable :: supercell_mask real, dimension(:, :), allocatable :: theta_2d complex, dimension(:, :), allocatable :: extracted_field integer :: isc, ncell, it0, icell integer :: nextend, iend, istart, ntheta_minus_1 logical :: drop_duplicate_theta_points ! Get the supercell label related to the point of interest and ! form a logical mask selecting members of this supercell isc = supercell_labels(it, ik) allocate(supercell_mask(ntheta0)) supercell_mask = supercell_labels(:, ik) == isc ! Count how many cells there are. ncell = count(supercell_mask) theta0_values = pack(theta0(:, ik), supercell_mask) theta0_indices_real = pack([(it0*1.0, it0 = 1, ntheta0)], supercell_mask) ! Here we're sorting theta0 indices for member of the supercell ! with the theta0 values as a key so that theta0_indices_real will give ! the it values of the supercell members in order of increasing theta0. ! We could skip this sorting and just work out where we transition from the ! +ve kx to the -ve kx and slice there, but sorting should work in general ! and avoids potentially confusing integer index calculations call quicksort(ncell, theta0_values, theta0_indices_real) ! Convert theta0_indices_real to integer to allow use for indexing ! arrays later. allocate(theta0_indices(size(theta0_indices_real))) theta0_indices = int(theta0_indices_real) theta0_indices = theta0_indices(ncell:1:-1) allocate(extracted_field(-ntgrid:ntgrid, ncell)) allocate(theta_2d(-ntgrid:ntgrid, ncell)) do icell = 1, ncell extracted_field(:, icell) = field_like_ft(:, theta0_indices(icell), ik) theta_2d(:, icell) = theta - theta0(theta0_indices(icell), ik) end do drop_duplicate_theta_points = get_option_with_default(drop_duplicates, .false.) if (drop_duplicate_theta_points) then ! Get the size of the domain, dropping one theta point for all but the last cell nextend = size(theta_2d) - (ncell - 1) allocate(theta_bs(nextend)) allocate(field_like_bs(nextend)) ntheta_minus_1 = size(theta) - 1 istart = 1 iend = istart + ntheta_minus_1 - 1 do icell = 1, ncell - 1 theta_bs(istart:iend) = theta_2d(:ntgrid-1, icell) field_like_bs(istart:iend) = extracted_field(:ntgrid-1, icell) istart = iend + 1 iend = istart + ntheta_minus_1 - 1 end do theta_bs(istart:) = theta_2d(:, ncell) field_like_bs(istart:) = extracted_field(:, ncell) else theta_bs = reshape(theta_2d, [size(theta_2d)]) field_like_bs = reshape(extracted_field, [size(extracted_field)]) end if end subroutine flux_tube_field_like_to_ballooning_space !> FIXME : Add documentation subroutine check_kt_grids(report_unit) use kt_grids_single, only: check_kt_grids_single use kt_grids_range, only: check_kt_grids_range use kt_grids_box, only: check_kt_grids_box use kt_grids_specified, only: check_kt_grids_specified implicit none integer, intent(in) :: report_unit select case (gridopt_switch) case (gridopt_single) call check_kt_grids_single (rhostar, report_unit) case (gridopt_range) call check_kt_grids_range (rhostar, report_unit) case (gridopt_specified) call check_kt_grids_specified (report_unit, aky, theta0(:,1)) case (gridopt_box) call check_kt_grids_box (rhostar, report_unit) end select end subroutine check_kt_grids #include "kt_grids_auto_gen.inc" end module kt_grids