add_nl Subroutine

  subroutine add_nl (g_in, g1, phi, apar, bpar, adjust)
    use theta_grid, only: ntgrid, kxfac
    use gs2_layouts, only: g_lo, ik_idx, it_idx
    use dist_fn_arrays, only: g_adjust, from_g_gs2
    use gs2_transforms, only: transform2, inverse2
    use kt_grids, only: aky, akx
    use gs2_time, only: save_dt_cfl, check_time_step_too_large
    use constants, only: zi
    use unit_tests, only: debug_message
    use optionals, only: get_option_with_default
    use array_utils, only: copy, gs2_max_abs
    implicit none
    complex, dimension (-ntgrid:,:,g_lo%llim_proc:), intent (in) :: g_in
    complex, dimension (-ntgrid:,:,g_lo%llim_proc:), intent (out) :: g1
    !> Local array used to store the adjusted input dfn if needed. Saves
    !> calling g_adjust twice.
    complex, dimension (-ntgrid:ntgrid, 2, g_lo%llim_proc:g_lo%ulim_alloc) :: g2
    complex, dimension (-ntgrid:,:,:), intent (in) :: phi, apar, bpar
    logical, intent(in), optional :: adjust
    logical :: should_adjust
    integer, parameter :: verb = 4
    integer :: iglo, ik, it
    real :: max_vel_x, max_vel_y
    
    call debug_message(verb, 'nonlinear_terms::add_nl starting')
    should_adjust = get_option_with_default(adjust, .true.)

    !Form g1=i*kx*chi
    call load_kx_phi(g1, phi)    
    call load_kx_bpar(g1, bpar)    
    call load_kx_apar(g1, apar)    

    call debug_message(verb, 'nonlinear_terms::add_nl calling transform2')

    !Transform to real space
    if (accelerated) then
       call transform2 (g1, aba)
       max_vel_y = gs2_max_abs(aba)
    else
       call transform2 (g1, ba)
       max_vel_y = gs2_max_abs(ba)
    end if
    max_vel_y = max_vel_y * cfly

    call debug_message(verb, 'nonlinear_terms::add_nl calculated dphi/dx')
    
    !Form g1=i*ky*g_wesson
    call copy(g_in, g1)
    if (should_adjust) call g_adjust(g1, phi, bpar, direction = from_g_gs2)
    call copy(g1, g2)

    !$OMP PARALLEL DO DEFAULT(none) &
    !$OMP PRIVATE(iglo, ik) &
    !$OMP SHARED(g_lo, g1, aky) &
    !$OMP SCHEDULE(static)
    do iglo = g_lo%llim_proc, g_lo%ulim_proc
       ik = ik_idx(g_lo,iglo)
       g1(:,:,iglo)=g1(:,:,iglo)*zi*aky(ik)
    enddo
    !$OMP END PARALLEL DO

    !Transform to real space
    if (accelerated) then
       call transform2 (g1, agb)
    else
       call transform2 (g1, gb)
    end if

    call debug_message(verb, 'nonlinear_terms::add_nl calculated dg/dy')

    call calculate_bracket(.true.)
    
    call debug_message(verb, 'nonlinear_terms::add_nl calculated dphi/dx.dg/dy')

    !Form g1=i*ky*chi
    call load_ky_phi(g1, phi)
    call load_ky_bpar(g1, bpar)
    call load_ky_apar(g1, apar)

    call debug_message(verb, 'nonlinear_terms::add_nl calculated dphi/dy')

    !Transform to real space    
    if (accelerated) then
       call transform2 (g1, aba)
       max_vel_x = gs2_max_abs(aba)
    else
       call transform2 (g1, ba)
       max_vel_x = gs2_max_abs(ba)
    end if
    max_vel_x = max_vel_x * cflx

    !Form g1=i*kx*g_wesson noting that g2 holds our starting g_wesson

    !$OMP PARALLEL DO DEFAULT(none) &
    !$OMP PRIVATE(iglo, it) &
    !$OMP SHARED(g_lo, g2, akx) &
    !$OMP SCHEDULE(static)
    do iglo = g_lo%llim_proc, g_lo%ulim_proc
       it = it_idx(g_lo,iglo)
       g2(:,:,iglo) = g2(:,:,iglo) * zi * akx(it)
    enddo
    !$OMP END PARALLEL DO

    !Transform to real space
    if (accelerated) then
       call transform2 (g2, agb)
    else
       call transform2 (g2, gb)
    end if

    call debug_message(verb, 'nonlinear_terms::add_nl calculated dg/dx')

    call calculate_bracket(.false.)
    
    call debug_message(verb, &
      'nonlinear_terms::add_nl calculated dphi/dx.dg/dy.dphi/dy.dg/dx')

    !Transform NL source back to spectral space
    if (accelerated) then
       call inverse2 (abracket, g1)
    else
       call inverse2 (bracket, g1)
    end if

    !Store local max vel components for future max reduce
    max_vel_components(1) = max_vel_x
    max_vel_components(2) = max_vel_y

    ! Scale by kxfac
    !$OMP PARALLEL DO DEFAULT(none) &
    !$OMP PRIVATE(iglo) &
    !$OMP SHARED(g_lo, g1, kxfac) &
    !$OMP SCHEDULE(static)
    do iglo = g_lo%llim_proc, g_lo%ulim_proc
       g1(:, :, iglo) = g1(:, :, iglo) * kxfac
    end do
    !$OMP END PARALLEL DO

    call debug_message(verb, &
      'nonlinear_terms::add_nl calculated nonlinear term')

  end subroutine add_nl