neoclassical_terms.f90 Source File

module neoclassical_terms

  use debug_flags, only: debug => neoclassical_terms_debug
   implicit none

   public :: init_neoclassical_terms
   public :: finish_neoclassical_terms
   public :: include_neoclassical_terms
   public :: dfneo_dzed, dfneo_dvpa, dfneo_drho, dfneo_dalpha
   public :: dphineo_dzed, dphineo_drho, dphineo_dalpha

   private

   logical :: include_neoclassical_terms
   integer :: nradii
   real :: drho

   integer :: neo_option_switch
   integer, parameter :: neo_option_sfincs = 1

   real, dimension(:, :, :), allocatable :: dfneo_dzed, dfneo_dvpa, dfneo_drho, dfneo_dalpha
   real, dimension(:, :), allocatable :: dphineo_dzed, dphineo_drho, dphineo_dalpha

   logical :: neoinit = .false.

contains

   subroutine init_neoclassical_terms

      use mp, only: proc0 
      use zgrid, only: nzgrid
      use parameters_kxky_grids, only: nalpha
      use vpamu_grids, only: nvpa, nmu
      use species, only: nspec
      use stella_layouts, only: vmu_lo
      use sfincs_interface, only: get_neo_from_sfincs

      implicit none

      real, dimension(:, :, :, :, :, :), allocatable :: f_neoclassical
      real, dimension(:, :, :), allocatable :: phi_neoclassical
      real, dimension(:, :, :, :, :), allocatable :: dfneo_dalpha_local

      integer :: iz, ialpha

      if (neoinit) return
      neoinit = .true.

      debug = debug .and. proc0
      
      call read_parameters
      if (include_neoclassical_terms) then
         allocate (f_neoclassical(nalpha, -nzgrid:nzgrid, nvpa, nmu, nspec, -nradii / 2:nradii / 2))
         allocate (phi_neoclassical(nalpha, -nzgrid:nzgrid, -nradii / 2:nradii / 2))
         allocate (dfneo_dalpha_local(nalpha, -nzgrid:nzgrid, nvpa, nmu, nspec))
         if (.not. allocated(dfneo_dvpa)) &
            allocate (dfneo_dvpa(nalpha, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
         if (.not. allocated(dfneo_drho)) &
            allocate (dfneo_drho(nalpha, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
         if (.not. allocated(dfneo_dzed)) &
            allocate (dfneo_dzed(nalpha, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
         if (.not. allocated(dfneo_dalpha)) &
            allocate (dfneo_dalpha(nalpha, -nzgrid:nzgrid, vmu_lo%llim_proc:vmu_lo%ulim_alloc))
         if (.not. allocated(dphineo_dzed)) &
            allocate (dphineo_dzed(nalpha, -nzgrid:nzgrid))
         if (.not. allocated(dphineo_drho)) &
            allocate (dphineo_drho(nalpha, -nzgrid:nzgrid))
         if (.not. allocated(dphineo_dalpha)) &
            allocate (dphineo_dalpha(nalpha, -nzgrid:nzgrid))
         select case (neo_option_switch)
         case (neo_option_sfincs)
            call get_neo_from_sfincs(nradii, drho, f_neoclassical, phi_neoclassical, dfneo_dalpha_local, dphineo_dalpha)
         end select
         if (debug) write (6, *) 'neoclassical_terms::init_neoclassical_terms::get_dfneo_dzed'
         call get_dfneo_dzed(f_neoclassical(:, :, :, :, :, 0), dfneo_dzed)
         if (debug) write (6, *) 'neoclassical_terms::init_neoclassical_terms::get_dfneo_dvpa'
         call get_dfneo_dvpa(f_neoclassical(:, :, :, :, :, 0), dfneo_dvpa)
         if (debug) write (6, *) 'neoclassical_terms::init_neoclassical_terms::get_dfneo_drho'
         call get_dfneo_drho(f_neoclassical, dfneo_drho)
         if (debug) write (6, *) 'neoclassical_terms::init_neoclassical_terms::get_dphineo_dzed'
         call get_dphineo_dzed(phi_neoclassical(:, :, 0), dphineo_dzed)
         if (debug) write (6, *) 'neoclassical_terms::init_neoclassical_terms::get_dphineo_drho'
         call get_dphineo_drho(phi_neoclassical, dphineo_drho)
         do iz = -nzgrid, nzgrid
            do ialpha = 1, nalpha
               call distribute_vmus_over_procs(dfneo_dalpha_local(ialpha, iz, :, :, :), dfneo_dalpha(ialpha, iz, :))
            end do
         end do
         if (debug) write (6, *) 'neoclassical_terms::init_neoclassical_terms::write_neoclassical'
         call write_neoclassical(f_neoclassical, phi_neoclassical)
         deallocate (f_neoclassical, phi_neoclassical, dfneo_dalpha_local)
      end if

   end subroutine init_neoclassical_terms

   subroutine read_parameters

      use mp, only: proc0, broadcast
      use file_utils, only: error_unit, input_unit_exist
      use text_options, only: text_option, get_option_value

      implicit none

      type(text_option), dimension(2), parameter :: neoopts = (/ &
                                                    text_option('default', neo_option_sfincs), &
                                                    text_option('sfincs', neo_option_sfincs)/)
      character(10) :: neo_option

      namelist /neoclassical_input/ include_neoclassical_terms, &
         neo_option, nradii, drho

      logical :: exist
      integer :: ierr, in_file

      if (proc0) then
         ! set to .true. to include neoclassical terms in GK equation
         include_neoclassical_terms = .false.
         ! number of radial points used for radial derivatives
         ! of neoclassical quantities
         nradii = 5
         ! spacing in rhoc between points used for radial derivatives
         drho = 0.01
         ! option for obtaining neoclassical distribution function and potential
         neo_option = 'sfincs'

         in_file = input_unit_exist("neoclassical_input", exist)
         if (exist) read (unit=in_file, nml=neoclassical_input)

         ierr = error_unit()
         call get_option_value &
            (neo_option, neoopts, neo_option_switch, &
             ierr, "neo_option in neoclassical_input")

         if (nradii /= 3 .and. nradii /= 5) then
            write (*, *) 'WARNING: only nradii of 3 or 5 is currently supported in neoclassical_input namelist'
            write (*, *) 'WARNING: forcing nradii=5'
            nradii = 5
         end if
      end if

      call broadcast(include_neoclassical_terms)
      call broadcast(neo_option_switch)
      call broadcast(nradii)
      call broadcast(drho)

   end subroutine read_parameters

   subroutine distribute_vmus_over_procs(local, distributed)

      use stella_layouts, only: vmu_lo
      use stella_layouts, only: iv_idx, imu_idx, is_idx

      implicit none

      real, dimension(:, :, :), intent(in) :: local
      real, dimension(vmu_lo%llim_proc:), intent(out) :: distributed

      integer :: ivmu, iv, imu, is

      do ivmu = vmu_lo%llim_proc, vmu_lo%ulim_proc
         iv = iv_idx(vmu_lo, ivmu)
         imu = imu_idx(vmu_lo, ivmu)
         is = is_idx(vmu_lo, ivmu)
         distributed(ivmu) = local(iv, imu, is)
      end do

   end subroutine distribute_vmus_over_procs

   subroutine get_dfneo_dvpa(fneo, dfneo)

      use finite_differences, only: fd5pt
      use stella_layouts, only: vmu_lo
      use zgrid, only: nzgrid
      use vpamu_grids, only: nvpa, nmu
      use vpamu_grids, only: dvpa
      use species, only: nspec
      use parameters_kxky_grids, only: nalpha

      implicit none

      real, dimension(:, -nzgrid:, :, :, :), intent(in) :: fneo
      real, dimension(:, -nzgrid:, vmu_lo%llim_proc:), intent(out) :: dfneo

      integer :: ia, iz, imu, is
      real, dimension(:), allocatable :: tmp1, tmp2
      real, dimension(:, :, :, :, :), allocatable :: dfneo_local

      allocate (tmp1(nvpa), tmp2(nvpa))
      allocate (dfneo_local(nalpha, -nzgrid:nzgrid, nvpa, nmu, nspec))

      do is = 1, nspec
         do imu = 1, nmu
            do iz = -nzgrid, nzgrid
               do ia = 1, nalpha
                  ! hack to avoid dealing with negative indices in fd5pt
                  tmp1 = fneo(ia, iz, :, imu, is)
                  call fd5pt(tmp1, tmp2, dvpa)
                  dfneo_local(ia, iz, :, imu, is) = tmp2
               end do
            end do
         end do
      end do

      do iz = -nzgrid, nzgrid
         do ia = 1, nalpha
            call distribute_vmus_over_procs(dfneo_local(ia, iz, :, :, :), dfneo(ia, iz, :))
         end do
      end do

      deallocate (dfneo_local)
      deallocate (tmp1, tmp2)

   end subroutine get_dfneo_dvpa

   subroutine get_dfneo_dzed(fneo, dfneo)

      use finite_differences, only: fd5pt
      use zgrid, only: nztot, nzgrid, delzed
      use vpamu_grids, only: nvpa, nmu
      use species, only: nspec
      use stella_layouts, only: vmu_lo
      use parameters_kxky_grids, only: nalpha

      implicit none

      real, dimension(:, -nzgrid:, :, :, :), intent(in) :: fneo
      real, dimension(:, -nzgrid:, vmu_lo%llim_proc:), intent(out) :: dfneo

      integer :: iv, imu, is, iz, ia
      real, dimension(:), allocatable :: tmp1, tmp2
      real, dimension(:), allocatable :: dfneo_local(:, :, :, :, :)

      allocate (tmp1(nztot), tmp2(nztot))
      allocate (dfneo_local(nalpha, -nzgrid:nzgrid, nvpa, nmu, nspec))

      do is = 1, nspec
         do imu = 1, nmu
            do iv = 1, nvpa
               do ia = 1, nalpha
                  ! hack to avoid dealing with negative indices in fd5pt
                  tmp1 = fneo(ia, :, iv, imu, is)
                  call fd5pt(tmp1, tmp2, delzed(0))
                  dfneo_local(ia, :, iv, imu, is) = tmp2
               end do
            end do
         end do
      end do

      do iz = -nzgrid, nzgrid
         do ia = 1, nalpha
            call distribute_vmus_over_procs(dfneo_local(ia, iz, :, :, :), dfneo(ia, iz, :))
         end do
      end do

      deallocate (dfneo_local)
      deallocate (tmp1, tmp2)

   end subroutine get_dfneo_dzed

   subroutine get_dfneo_drho(fneo, dfneo)

      use finite_differences, only: fd3pt, fd5pt
      use zgrid, only: nzgrid
      use vpamu_grids, only: nvpa, nmu
      use species, only: nspec
      use stella_layouts, only: vmu_lo
      use parameters_kxky_grids, only: nalpha

      implicit none

      real, dimension(:, -nzgrid:, :, :, :, -nradii/2:), intent(in) :: fneo
      real, dimension(:, -nzgrid:, vmu_lo%llim_proc:), intent(out) :: dfneo

      integer :: ia, iz, iv, imu, is
      real, dimension(:), allocatable :: tmp1, tmp2
      real, dimension(:, :, :, :, :), allocatable :: dfneo_local

      allocate (tmp1(nradii), tmp2(nradii))
      allocate (dfneo_local(nalpha, -nzgrid:nzgrid, nvpa, nmu, nspec))

      do is = 1, nspec
         do imu = 1, nmu
            do iv = 1, nvpa
               do iz = -nzgrid, nzgrid
                  do ia = 1, nalpha
                     ! hack to avoid dealing with negative indices in fd5pt
                     tmp1 = fneo(ia, iz, iv, imu, is, :)
                     if (nradii == 5) then
                        call fd5pt(tmp1, tmp2, drho)
                     else
                        call fd3pt(tmp1, tmp2, drho)
                     end if
                     dfneo_local(ia, iz, iv, imu, is) = tmp2(nradii / 2 + 1)
                  end do
               end do
            end do
         end do
      end do

      do iz = -nzgrid, nzgrid
         do ia = 1, nalpha
            call distribute_vmus_over_procs(dfneo_local(ia, iz, :, :, :), dfneo(ia, iz, :))
         end do
      end do

      deallocate (dfneo_local)
      deallocate (tmp1, tmp2)

   end subroutine get_dfneo_drho

   subroutine get_dphineo_dzed(phineo, dphineo)

      use finite_differences, only: fd5pt
      use zgrid, only: nztot, nzgrid, delzed
      use parameters_kxky_grids, only: nalpha

      implicit none

      real, dimension(:, -nzgrid:), intent(in) :: phineo
      real, dimension(:, -nzgrid:), intent(out) :: dphineo

      integer :: ia
      real, dimension(:), allocatable :: tmp1, tmp2

      allocate (tmp1(nztot), tmp2(nztot))

      do ia = 1, nalpha
         ! hack to avoid dealing with negative indices in fd5pt
         tmp1 = phineo(ia, :)
         call fd5pt(tmp1, tmp2, delzed(0))
         dphineo(ia, :) = tmp2
      end do

      deallocate (tmp1, tmp2)

   end subroutine get_dphineo_dzed

   subroutine get_dphineo_drho(phineo, dphineo)

      use finite_differences, only: fd3pt, fd5pt
      use zgrid, only: nzgrid
      use parameters_kxky_grids, only: nalpha

      implicit none

      real, dimension(:, -nzgrid:, -nradii/2:), intent(in) :: phineo
      real, dimension(:, -nzgrid:), intent(out) :: dphineo

      integer :: iz, ia
      real, dimension(:), allocatable :: tmp1, tmp2

      allocate (tmp1(nradii), tmp2(nradii))

      do iz = -nzgrid, nzgrid
         do ia = 1, nalpha
            ! hack to avoid dealing with negative indices in fd5pt
            tmp1 = phineo(ia, iz, :)
            if (nradii == 5) then
               call fd5pt(tmp1, tmp2, drho)
            else
               call fd3pt(tmp1, tmp2, drho)
            end if
            dphineo(ia, iz) = tmp2(nradii / 2 + 1)
         end do
      end do

      deallocate (tmp1, tmp2)

   end subroutine get_dphineo_drho

   subroutine write_neoclassical(fnc, phinc)

      use mp, only: proc0
      use mp, only: send, receive
      use file_utils, only: open_output_file, close_output_file
      use zgrid, only: nzgrid, zed
      use vpamu_grids, only: vpa, mu
      use stella_layouts, only: vmu_lo
      use stella_layouts, only: iv_idx, imu_idx, is_idx
      use stella_layouts, only: idx_local, proc_id
      use parameters_kxky_grids, only: nalpha

      implicit none

      real, dimension(:, -nzgrid:, :, :, :, -nradii/2:), intent(in) :: fnc
      real, dimension(:, -nzgrid:, -nradii/2:), intent(in) :: phinc

      integer :: neo_unit
      integer :: irad, iz, ivmu, iv, imu, is, ia
      real, dimension(:, :), allocatable :: dfdv_local, dfdr_local, dfdz_local

      allocate (dfdv_local(nalpha, -nzgrid:nzgrid))
      allocate (dfdr_local(nalpha, -nzgrid:nzgrid))
      allocate (dfdz_local(nalpha, -nzgrid:nzgrid))

      if (proc0) then
         call open_output_file(neo_unit, '.neoclassical')
         write (neo_unit, '(3a8,10a13)') '#1.rad', '2.spec', '3.alpha', '4.zed', '5.mu', &
            '6.vpa', '7.f_neo', '8.dfdvpa_neo', '9.dfdrho_neo', '10.dfdzed_neo', '11.phi_neo', &
            '12.dphidrho', '13.dphidzed'
      end if
      do irad = -nradii / 2, nradii / 2
         do ivmu = vmu_lo%llim_world, vmu_lo%ulim_world
            iv = iv_idx(vmu_lo, ivmu)
            imu = imu_idx(vmu_lo, ivmu)
            is = is_idx(vmu_lo, ivmu)
            if (idx_local(vmu_lo, iv, imu, is)) then
               if (proc0) then
                  dfdv_local = dfneo_dvpa(:, :, ivmu)
                  dfdr_local = dfneo_drho(:, :, ivmu)
                  dfdz_local = dfneo_dzed(:, :, ivmu)
               else
                  call send(dfneo_dvpa(:, :, ivmu), 0)
                  call send(dfneo_drho(:, :, ivmu), 0)
                  call send(dfneo_dzed(:, :, ivmu), 0)
               end if
            else if (proc0) then
               call receive(dfdv_local, proc_id(vmu_lo, ivmu))
               call receive(dfdr_local, proc_id(vmu_lo, ivmu))
               call receive(dfdz_local, proc_id(vmu_lo, ivmu))
            end if
            if (proc0) then
               do iz = -nzgrid, nzgrid
                  do ia = 1, nalpha
                     write (neo_unit, '(3i8,10e13.5)') irad, is, ia, zed(iz), mu(imu), vpa(iv), &
                        fnc(ia, iz, iv, imu, is, irad), &
                        dfdv_local(ia, iz), &
                        dfdr_local(ia, iz), &
                        dfdz_local(ia, iz), &
                        phinc(ia, iz, irad), dphineo_drho(ia, iz), dphineo_dzed(ia, iz)
                  end do
               end do
            end if
         end do
         if (proc0) write (neo_unit, *)
      end do
      if (proc0) call close_output_file(neo_unit)

      deallocate (dfdv_local, dfdr_local, dfdz_local)

   end subroutine write_neoclassical

   subroutine finish_neoclassical_terms

      implicit none

      if (allocated(dfneo_dvpa)) deallocate (dfneo_dvpa)
      if (allocated(dfneo_drho)) deallocate (dfneo_drho)
      if (allocated(dfneo_dzed)) deallocate (dfneo_dzed)
      if (allocated(dfneo_dalpha)) deallocate (dfneo_dalpha)
      if (allocated(dphineo_dzed)) deallocate (dphineo_dzed)
      if (allocated(dphineo_drho)) deallocate (dphineo_drho)
      if (allocated(dphineo_dalpha)) deallocate (dphineo_dalpha)

      neoinit = .false.

   end subroutine finish_neoclassical_terms

end module neoclassical_terms