2078172070

Committed 03 Oct 2025 09:22AM UTC coverage: 10.835% (+0.03%) from 10.806%

Build # 2078172070

Build Type

push

gitlab-ci

Committed by David Dickinson

Commit Message

Merged in experimental/user_controlled_output_base_name (pull request #1184)

Run Details

5049 of 46599 relevant lines covered (10.83%)

119786.99 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/src/gs2_diagnostics.fpp

#include "unused_dummy.inc"
!> A module for calculating and writing the main GS2 outputs.
module gs2_diagnostics
  use gs2_heating, only: heating_diagnostics
  use gs2_save, only: save_many
  use diagnostics_configuration, only: diagnostics_config_type
  use gs2_io, only: get_netcdf_file_id

  implicit none

  private

  public :: read_parameters, init_gs2_diagnostics, finish_gs2_diagnostics
  public :: check_gs2_diagnostics, wnml_gs2_diagnostics, check_restart_file_writeable
  public :: loop_diagnostics, omegaavg, calculate_instantaneous_omega
  public :: diffusivity, calculate_simple_quasilinear_flux_metric_by_k
  public :: save_restart_dist_fn, do_write_fyx, do_write_f, do_write_parity
  public :: do_write_heating, do_write_nl_flux_dist, do_write_geom, heating
  public :: save_distfn, save_glo_info_and_grids, save_velocities, save_for_restart
  public :: omegatol, nwrite, nmovie, navg, nsave, make_movie, exit_when_converged
  public :: do_write_correlation_extend, do_write_correlation, do_write_symmetry
  public :: do_final_gs2_diagnostics

  real :: omegatinst, omegatol, damped_threshold, omega_checks_start
  logical :: print_line, print_flux_line, write_line, write_flux_line, write_collisional
  logical :: write_omega, write_omavg, write_ascii, write_gs, write_ql_metric, write_moments
  logical :: write_g, write_gyx, write_eigenfunc, write_fields, write_final_fields
  logical :: write_final_antot, write_final_moments, write_parity, write_final_db
  logical :: write_moments_over_time, write_fluxes_by_energy, write_energy_exchange
  logical :: write_moments_by_mode
  logical :: write_cross_phase, write_final_epar, write_kpar, write_heating, write_lorentzian
  logical :: write_fluxes, write_fluxes_by_mode, write_kinetic_energy_transfer, write_verr
  logical :: write_cerr, exit_when_converged, use_nonlin_convergence, write_zonal_transfer
  logical :: save_for_restart, save_distfn, save_glo_info_and_grids, save_velocities
  logical :: write_symmetry, write_correlation_extend, write_correlation, make_movie
  logical :: write_nl_flux_dist, file_safety_check, append_old, make_restart_dir
  logical :: write_phi_over_time, write_apar_over_time, write_bpar_over_time, write_jext
  integer :: nwrite, nmovie, navg, nsave, igomega, nwrite_mult, nc_sync_freq

  ! internal. `write_any` differs from new diagnostics input `write_any`. The
  ! latter can be used to turn off all output, whereas here, it's used to
  ! _check_ if any outputs are enabled
  logical :: write_any, write_any_fluxes
  logical, private :: initialized = .false.
  complex, allocatable, dimension (:,:,:) :: domega

  !> Complex frequency, and time-averaged complex frequency
  complex, dimension (:, :), allocatable :: omega, omegaavg

  integer :: out_unit, heat_unit, heat_unit2, lpc_unit
  integer :: jext_unit, phase_unit, dist_unit, yxdist_unit
  integer :: res_unit, res_unit2, parity_unit, cres_unit
  integer :: conv_nstep_av, conv_min_step, conv_max_step
  integer :: conv_nsteps_converged

  !> Frequency time history, size `(navg, ntheta0, naky)`
  complex, dimension (:,:,:), allocatable :: omegahist

  !> Heating diagnostics summed over space at the current timestep
  type (heating_diagnostics) :: h
  !> Heating diagnostics summed over space over the last [[gs2_diagnostics_knobs:navg]] timesteps
  type (heating_diagnostics), dimension(:), save, allocatable :: h_hist
  !> Heating diagnostics as a function of \((\theta, k_y)\) at the current timestep
  type (heating_diagnostics), dimension(:,:), save, allocatable :: hk
  !> Heating diagnostics as a function of \((\theta, k_y)\) over the last [[gs2_diagnostics_knobs:navg]] timesteps
  type (heating_diagnostics), dimension(:,:,:), save, allocatable :: hk_hist
  !GGH J_external
  real, dimension(:,:,:), allocatable ::  j_ext_hist

  real :: conv_test_multiplier

  integer :: nout = 1
  integer :: nout_movie = 1
  integer :: nout_big = 1
  complex :: wtmp_old = 0.

contains

  !> Write the diagnostics namelist to file
  subroutine wnml_gs2_diagnostics(unit)
    use diagnostics_configuration, only: diagnostics_config
    implicit none
    !> Unit of an open file to write to
    integer, intent(in) :: unit
    call diagnostics_config%write(unit)
  end subroutine wnml_gs2_diagnostics

  !> Perform some basic checking of the diagnostic input parameters, and write the results to file
  subroutine check_gs2_diagnostics(report_unit)
    use file_utils, only: run_name
    use nonlinear_terms, only: nonlin
    use dist_fn, only : def_parity, even 
    use init_g, only : restart_file
    use gs2_save, only: restart_writable
    implicit none
    !> Unit of an open file to write to
    integer, intent(in) :: report_unit
    logical :: writable
    write (report_unit, *) 
    write (report_unit, fmt="('------------------------------------------------------------')")
    write (report_unit, *) 
    write (report_unit, fmt="('Diagnostic control section.')")

    if (print_line) then
       write (report_unit, fmt="('print_line = T:            Estimated frequencies &
            & output to the screen every ',i4,' steps.')") nwrite
    else
       ! nothing
    end if

    if (write_line) then
       if (write_ascii) then
          write (report_unit, fmt="('write_line = T:            Estimated frequencies output to ',a,' every ',i4,' steps.')") &
               & trim(run_name)//'.out',  nwrite
       end if
       write (report_unit, fmt="('write_line = T:            Estimated frequencies output to ',a,' every ',i4,' steps.')") &
            & trim(run_name)//'.out.nc',  nwrite
    else
       ! nothing
    end if

    if (print_flux_line) then
       write (report_unit, fmt="('print_flux_line = T:       Instantaneous fluxes output to screen every ', &
            & i4,' steps.')") nwrite
    else
       ! nothing
    end if

    if (write_flux_line) then
       if (write_ascii) then
          write (report_unit, fmt="('write_flux_line = T:       Instantaneous fluxes output to ',a,' every ',i4,' steps.')") &
               & trim(run_name)//'.out',  nwrite
       end if
       write (report_unit, fmt="('write_flux_line = T:       Instantaneous fluxes output to ',a,' every ',i4,' steps.')") &
            & trim(run_name)//'.out.nc',  nwrite
    else
       ! nothing
    end if

    if (write_omega) then
       if (write_ascii) then
          write (report_unit, fmt="('write_omega = T:           Instantaneous frequencies written to ',a)") trim(run_name)//'.out'
       else
          write (report_unit, fmt="('write_omega = T:           No effect.')")
       end if
       write (report_unit, fmt="('                           Frequencies calculated at igomega = ',i4)") igomega
       if (def_parity .and. .not. even) then
          write (report_unit, fmt="('################# WARNING #######################')")
          write (report_unit, fmt="('   You probably want igomega /= 0 for odd parity modes.')") 
          write (report_unit, fmt="('################# WARNING #######################')")
          write (report_unit, *) 
       end if
    end if

    if (write_omavg) then
       if (write_ascii) then
          write (report_unit, fmt="('write_omavg = T:           Time-averaged frequencies written to ',a)") trim(run_name)//'.out'
          write (report_unit, fmt="('                           Averages taken over ',i4,' timesteps.')") navg
       else
          write (report_unit, fmt="('write_omavg = T:           No effect.')")
       end if
    end if

    if (write_ascii) then
       write (report_unit, fmt="('write_ascii = T:           Write some data to ',a)") trim(run_name)//'.out'
    end if

    if (write_eigenfunc) then
       if (write_ascii) then
          write (report_unit, fmt="('write_eigenfunc = T:       Normalized Phi(theta) written to ',a)") trim(run_name)//'.eigenfunc'
       end if
       write (report_unit, fmt="('write_eigenfunc = T:       Normalized Phi(theta) written to ',a)") trim(run_name)//'.out.nc'
    end if

    if (write_final_fields) then
       if (write_ascii) then
          write (report_unit, fmt="('write_final_fields = T:    Phi(theta), etc. written to ',a)") trim(run_name)//'.fields'
       end if
       write (report_unit, fmt="('write_final_fields = T:    Phi(theta), etc. written to ',a)") trim(run_name)//'.out.nc'
    end if

    if (write_final_antot) then
       if (write_ascii) then
          write (report_unit, fmt="('write_final_antot = T:          Sources for Maxwell eqns. written to ',a)") &
               & trim(run_name)//'.antot'
       end if
       write (report_unit, fmt="('write_final_antot = T:          Sources for Maxwell eqns. written to ',a)") &
            & trim(run_name)//'.out.nc'
    end if

    if (write_final_moments) then
       if (write_ascii) then
          write (report_unit, fmt="('write_final_moments = T:   Low-order moments of g written to ',a)") &
               & trim(run_name)//'.moments'
          write (report_unit, fmt="('write_final_moments = T:   int dl/B average of low-order moments of g written to ',a)") &
               & trim(run_name)//'.amoments'
       end if
       write (report_unit, fmt="('write_final_moments = T:   Low-order moments of g written to ',a)") &
            & trim(run_name)//'.out.nc'
       write (report_unit, fmt="('write_final_moments = T:   int dl/B average of low-order moments of g written to ',a)") &
            & trim(run_name)//'.out.nc'
    end if

    if (write_final_epar) then
       if (write_ascii) then
          write (report_unit, fmt="('write_final_epar = T:      E_parallel(theta) written to ',a)") trim(run_name)//'.epar'
       end if
       write (report_unit, fmt="('write_final_epar = T:      E_parallel(theta) written to ',a)") trim(run_name)//'.out.nc'
    end if

    if (write_fluxes) then
       if (write_ascii) then
          write (report_unit, fmt="('write_fluxes = T:         Phi**2(kx, ky) written to ',a)") trim(run_name)//'.out'
       end if
    else
       write (report_unit, fmt="('write_fluxes = F:         Phi**2(kx, ky) NOT written to ',a)") trim(run_name)//'.out'
    end if

    if (save_for_restart) then
       write (report_unit, fmt="('save_for_restart = T:      Restart files written to ',a)") trim(restart_file)//'.(PE)'
    else
       if (nonlin) then
          write (report_unit, *) 
          write (report_unit, fmt="('################# WARNING #######################')")
          write (report_unit, fmt="('save_for_restart = F:              This run cannot be continued.')")
          write (report_unit, fmt="('THIS IS PROBABLY AN ERROR.')") 
          write (report_unit, fmt="('################# WARNING #######################')")
          write (report_unit, *) 
       end if
    end if

    !Verify restart file can be written
    if((save_for_restart.or.save_distfn).and.(file_safety_check))then
       !Can we write file?
       writable=restart_writable()

       !If we can't write the restart file then we should probably quit
       if((.not.writable))then
          if(save_for_restart)then 
             write (report_unit, *) 
             write (report_unit, fmt="('################# WARNING #######################')")
             write (report_unit, fmt="('save_for_restart = T:   But we cannot write to a test file like ',A,'.')") trim(restart_file)
             write (report_unit, fmt="('THIS IS PROBABLY AN ERROR --> Check restart_dir.')") 
             write (report_unit, fmt="('################# WARNING #######################')")
             write (report_unit, *) 
          endif
          if(save_distfn)then 
             write (report_unit, *) 
             write (report_unit, fmt="('################# WARNING #######################')")
             write (report_unit, fmt="('save_distfn = T:   But we cannot write to a test file like ',A,'.')") trim(restart_file)
             write (report_unit, fmt="('THIS IS PROBABLY AN ERROR --> Check restart_dir.')") 
             write (report_unit, fmt="('################# WARNING #######################')")
             write (report_unit, *) 
          endif
       endif
    endif

  end subroutine check_gs2_diagnostics

  !> Initialise this module and all its dependencies, including defining NetCDF
  !> vars, call real_init, which calls read_parameters; broadcast all the
  !> different write flags.
  subroutine init_gs2_diagnostics (gs2_diagnostics_config_in, header)
    use kt_grids, only: ntheta0, naky
    use species, only: nspec
    use gs2_io, only: init_gs2_io
    use gs2_heating, only: init_htype
    use collisions, only: collision_model_switch, init_lorentz_error
    use collisions, only: collision_model_lorentz, collision_model_lorentz_test
    use mp, only: broadcast, mp_abort
    use diagnostics_final_routines, only: init_par_filter
    use diagnostics_velocity_space, only: init_weights, setup_legendre_transform
    use standard_header, only: standard_header_type
    implicit none
    !> Input parameters for this module
    type(diagnostics_config_type), intent(in), optional :: gs2_diagnostics_config_in
    !> Header for files with build and run information
    type(standard_header_type), intent(in), optional :: header
    ! Actual value for optional `header` input
    type(standard_header_type) :: local_header

    if (present(header)) then
      local_header = header
    else
      local_header = standard_header_type()
    end if

    if (initialized) return
    initialized = .true.

    call real_init (gs2_diagnostics_config_in, header=local_header)

    ! initialize weights for less accurate integrals used
    ! to provide an error estimate for v-space integrals (energy and untrapped)
    if (write_verr) then
       call init_weights
       call setup_legendre_transform
    end if

    ! allocate heating diagnostic data structures
    if (write_heating) then
       allocate (h_hist(0:navg-1))
       call init_htype (h_hist,  nspec)

       allocate (hk_hist(ntheta0,naky,0:navg-1))
       call init_htype (hk_hist, nspec)

       call init_htype (h,  nspec)

       allocate (hk(ntheta0, naky))
       call init_htype (hk, nspec)
    else
       allocate (h_hist(0))
       allocate (hk(1,1))
       allocate (hk_hist(1,1,0))
    end if

    !GGH Allocate density and velocity perturbation diagnostic structures
    if (write_jext) then
       allocate (j_ext_hist(ntheta0, naky,0:navg-1))
       j_ext_hist = 0
    end if

    call init_gs2_io (append_old, make_movie, nout, header=local_header)

    call broadcast (nout)

    if (write_cerr) then
       if (collision_model_switch == collision_model_lorentz .or. &
            collision_model_switch == collision_model_lorentz_test) then
          call init_lorentz_error
       else
          write_cerr = .false.
       end if
    end if

    call check_restart_file_writeable(file_safety_check, save_for_restart, save_distfn, &
         make_restart_dir)

    !Setup the parallel fft if we're writing/using the parallel spectrum
    if (write_kpar .or. write_gs) call init_par_filter
  end subroutine init_gs2_diagnostics

  !> Check if we can write the restart files
  !>
  !> If restart files aren't writable, aborts if `need_restart` is
  !> true, or sets `extra_files` to false. Also prints a warning in
  !> the second case.
  subroutine check_restart_file_writeable(check_writable, need_restart, extra_files, &
       create_restart_dir)
    use init_g, only: get_restart_dir
    use gs2_save, only: restart_writable
    use mp, only: proc0, mp_abort, broadcast
    use file_utils, only: mkdir
    implicit none
    !> Has the user requested this check
    logical, intent(in) :: check_writable
    !> Has the user requested restart files
    logical, intent(in) :: need_restart
    !> Has the user requested distribution function to be written
    logical, intent(inout) :: extra_files
    !> Do we want to try to create the restart dir if initial writes fail?
    logical, intent(inout) :: create_restart_dir

    ! Error message from trying to open restart file
    character(len=:), allocatable :: error_message
    ! GS2 error message to show user
    character(len=:), allocatable :: abort_message

    ! Assume everything's fine if we're not checking, or if it's not needed
    if (.not. check_writable) return
    if (.not. (need_restart .or. extra_files)) return

    ! If we can write the restart files, we're done
    if (restart_writable(error_message=error_message)) return

    if (create_restart_dir) then
       if (proc0) then
          write(*, '(A," (",A,") ",A)') "Initial attempt to write to restart_dir", trim(get_restart_dir()),"failed, attempting to create it."
          call mkdir(get_restart_dir())
       end if

       ! If we can now write the restart files, we're done
       if (restart_writable(error_message=error_message)) return
    end if

    abort_message = " Cannot write restart files! Error message was:" // new_line('a') &
         // "    " // error_message // new_line('a') &
         // " Please check that `init_g_knobs::restart_dir = " // trim(get_restart_dir()) &
         // "` exists and has the correct permissions."

    ! User has requested restart files, but we can't write them => abort
    if (need_restart) then
      abort_message = abort_message // new_line('a') // "    --> Aborting"
      call mp_abort(trim(abort_message), to_screen=.true.)
    end if

    ! If it's just a case of save_distfn, then we can carry on but
    ! disable `save_distfn` and print a useful mesasge
    if (extra_files) then
      if(proc0) write(*, '(a, a, a)') abort_message, new_line('a'), "    --> Setting `save_distfn = F`"
      extra_files = .false.
    endif
  end subroutine check_restart_file_writeable

  !> Read the input parameters; open the various text output files according to
  !> the relevant input flags; allocate and zero the module-level flux arrays,
  !> [[gs2_diagnostics_knobs:omega]] and [[gs2_diagnostics_knobs:omegaavg]]
  subroutine real_init(gs2_diagnostics_config_in, header)
    use run_parameters, only: has_apar
    use file_utils, only: open_output_file
    use kt_grids, only: naky, ntheta0
    use mp, only: proc0
    use standard_header, only: standard_header_type
    use diagnostics_fluxes, only: init_diagnostics_fluxes
    use diagnostics_kinetic_energy_transfer, only: init_diagnostics_kinetic_energy_transfer
    use diagnostics_nonlinear_convergence, only: init_nonlinear_convergence
    use collisional_heating, only: init_collisional
    use diagnostics_zonal_transfer, only: init_diagnostics_transfer
    implicit none
    !> Input parameters for this module
    type(diagnostics_config_type), intent(in), optional :: gs2_diagnostics_config_in
    !> Header for files with build and run information
    type(standard_header_type), intent(in), optional :: header
    ! Actual value for optional `header` input
    type(standard_header_type) :: local_header

    if (present(header)) then
      local_header = header
    else
      local_header = standard_header_type()
    end if

    call read_parameters(gs2_diagnostics_config_in)

    if (proc0) then
       if (write_ascii) then
          call open_output_file (out_unit, ".out")          
       end if

       if (write_cross_phase .and. write_ascii) then
          call open_output_file (phase_unit, ".phase")
       end if

       if (write_heating .and. write_ascii) then
          call open_output_file (heat_unit, ".heat")
          call open_output_file (heat_unit2, ".heat2")
       end if

       if (write_verr .and. write_ascii) then
          call open_output_file (res_unit, ".vres")
          call open_output_file (lpc_unit, ".lpc")
          call open_output_file (res_unit2, ".vres2")
       end if

       if (write_cerr .and. write_ascii) then
          call open_output_file (cres_unit, ".cres")
       end if

       if (write_parity .and. write_ascii) then
          call open_output_file (parity_unit, ".parity")
       end if

       if (write_g .and. write_ascii) then
          call open_output_file(dist_unit, ".dist")
       end if

       if (write_gyx .and. write_ascii) then
          call open_output_file(yxdist_unit, ".yxdist")
       end if

       !GGH J_external, only if A_parallel is being calculated.
       if (write_jext .and. has_apar) then
          if (write_ascii) then
             call open_output_file (jext_unit, ".jext")
          end if
       else
          write_jext = .false.
       end if

       if (write_ascii) then
          write (unit=out_unit, fmt='(a)') local_header%to_string(file_description="GS2 output file")
       end if

       allocate (omegahist(0:navg-1,ntheta0,naky))
       omegahist = 0.0
       if (use_nonlin_convergence) call init_nonlinear_convergence(conv_nstep_av, nwrite)

    end if

    ! Needed to allow us to use common_calculate_fluxes later
    call init_diagnostics_fluxes()
    if (write_collisional) call init_collisional
    if (write_zonal_transfer) call init_diagnostics_transfer
    if (write_kinetic_energy_transfer) call init_diagnostics_kinetic_energy_transfer
    if (.not. allocated(omega)) then
       allocate(omega(ntheta0, naky))
       allocate(omegaavg(ntheta0, naky))
       omega=0
       omegaavg=0
    end if

  end subroutine real_init

  !> Read the input parameters for the diagnostics module
  subroutine read_parameters (gs2_diagnostics_config_in, warn_nonfunctional)
    use diagnostics_configuration, only: warn_about_nonfunctional_selection, diagnostics_config, check_and_disable_flags
    use optionals, only: get_option_with_default
    implicit none
    !> Configuration for this module, can be used to set new default values or
    !> avoid reading the input file
    type(diagnostics_config_type), intent(in), optional :: gs2_diagnostics_config_in
    logical, intent(in), optional :: warn_nonfunctional
    if (present(gs2_diagnostics_config_in)) diagnostics_config = gs2_diagnostics_config_in

    call diagnostics_config%init(name = 'gs2_diagnostics_knobs', requires_index = .false.)

    ! Print some health warnings if switches are not their default
    ! values and are not available in this diagnostics module
    if (get_option_with_default(warn_nonfunctional, .true.)) then
       call warn_about_nonfunctional_selection(.not. diagnostics_config%write_any, "write_any")
    end if

    ! Turn off diagnostics which aren't sensible with setup
    call check_and_disable_flags(diagnostics_config)

    ! Copy out internal values into module level parameters
    associate(self => diagnostics_config)
#include "diagnostics_copy_out_auto_gen.inc"
    end associate

    write_any = write_line .or. write_omega     .or. write_omavg &
         .or. write_flux_line .or. write_fluxes  &
         .or. write_kpar   .or. write_heating     .or. write_lorentzian  .or. write_gs
    write_any_fluxes = write_flux_line .or. print_flux_line .or. write_fluxes .or. write_fluxes_by_mode
  end subroutine read_parameters

  subroutine do_final_gs2_diagnostics
    use mp, only: proc0
    use fields_arrays, only: phinew, bparnew
    use antenna, only: dump_ant_amp
    use gs2_time, only: user_time
    use unit_tests, only: debug_message
    use diagnostics_final_routines, only: do_write_final_fields, do_write_kpar
    use diagnostics_final_routines, only: do_write_final_epar, do_write_final_db
    use diagnostics_final_routines, only: do_write_final_antot, do_write_final_moments
    use diagnostics_final_routines, only: do_write_gs
    use diagnostics_velocity_space, only: collision_error
    use diagnostics_fields, only: get_phi0
    implicit none

    call debug_message(3, 'gs2_diagnostics::final_gs2_diagnostics &
         & starting')

    if (write_gyx) call do_write_fyx (yxdist_unit, phinew, bparnew)
    if (write_g) call do_write_f (dist_unit)
    if (write_cerr) call collision_error (cres_unit, phinew, bparnew)

    if (proc0) then
       if (write_eigenfunc) call do_write_eigenfunc(get_netcdf_file_id(), write_ascii)
       if (write_final_fields) call do_write_final_fields(write_ascii, &
            file_id=get_netcdf_file_id())
       if (write_kpar) call do_write_kpar(write_ascii)
       if (write_final_epar) call do_write_final_epar(write_ascii, &
            file_id=get_netcdf_file_id())

       ! definition here assumes we are not using wstar_units
       if (write_final_db) call do_write_final_db(write_ascii)
    end if

    !Note pass in phase factor phi0 which may not be initialised
    !this is ok as phi0 will be set in routine if not already set
    if (write_final_moments) call do_write_final_moments(get_phi0(), &
         use_normalisation=write_eigenfunc, write_text=write_ascii, &
         file_id=get_netcdf_file_id())

    if (write_final_antot) call do_write_final_antot(write_ascii, &
         file_id=get_netcdf_file_id())

    call save_restart_dist_fn(save_for_restart, save_distfn, &
                              save_glo_info_and_grids=save_glo_info_and_grids, &
                              save_velocities=save_velocities, &
                              user_time=user_time)

    if (proc0) call dump_ant_amp

    if (write_gs) call do_write_gs(write_ascii)

    if (proc0) call do_write_geom
  end subroutine do_final_gs2_diagnostics

  !> Finalise the diagnostics module, writing final timestep diagnostics and
  !> closing any open files
  subroutine finish_gs2_diagnostics
    use gs2_io, only: nc_finish, get_netcdf_file_id
    use diagnostics_velocity_space, only: finish_legendre_transform, finish_weights
    if (.not. initialized) return
    ! Close some of the open ascii output files
    call close_files
    !Finalise the netcdf file
    call nc_finish(get_netcdf_file_id())
    !Now tidy up
    call deallocate_arrays
    if (write_verr) then
       call finish_weights
       call finish_legendre_transform
    end if
    wtmp_old = 0. ; nout = 1 ; nout_movie = 1 ; nout_big = 1
    initialized = .false.
  end subroutine finish_gs2_diagnostics

  !> Save some extra information in final restart files
  subroutine save_restart_dist_fn(save_for_restart, save_distfn, &
                                  save_glo_info_and_grids, &
                                  save_velocities, user_time, fileopt_base)
    use run_parameters, only: has_phi, has_apar, has_bpar
    use collisions, only: vnmult
    use gs2_save, only: gs2_save_for_restart
    use gs2_time, only: code_dt, code_dt_prev1, code_dt_prev2, code_dt_max
    use fields_arrays, only: phinew, bparnew
    use dist_fn_arrays, only: gnew, g_adjust, to_g_gs2, from_g_gs2
    use optionals, only: get_option_with_default
    !> See [[gs2_diagnostics_knobs:save_for_restart]]
    logical, intent(in) :: save_for_restart
    !> See [[gs2_diagnostics_knobs:save_distfn]]
    logical, intent(in) :: save_distfn
    !> See [[gs2_diagnostics_knobs:save_glo_info_and_grids]]
    logical, intent(in) :: save_glo_info_and_grids
    !> See [[gs2_diagnostics_knobs:save_velocities]]
    logical, intent(in) :: save_velocities
    !> Current simulation time
    real, intent(in) :: user_time
    !> Optional string to add to file name
    character(len=*), intent(in), optional :: fileopt_base
    character(len=:), allocatable :: fileopt

    fileopt = get_option_with_default(fileopt_base, "")

    if (save_for_restart) then
       call gs2_save_for_restart(gnew, user_time, vnmult, &
            has_phi, has_apar, has_bpar, &
            code_dt, code_dt_prev1, code_dt_prev2, code_dt_max, &
            save_glo_info_and_grids=save_glo_info_and_grids, &
            save_velocities=save_velocities, &
            fileopt = fileopt)
    end if

    if (save_distfn) then
       !Convert h to distribution function
       call g_adjust(gnew, phinew, bparnew, direction = from_g_gs2)

       !Save dfn, fields and velocity grids to file
       call gs2_save_for_restart(gnew, user_time, vnmult, &
            has_phi, has_apar, has_bpar, &
            code_dt, code_dt_prev1, code_dt_prev2, code_dt_max, &
            fileopt= fileopt//".dfn", &
            save_glo_info_and_grids=save_glo_info_and_grids, &
            save_velocities=save_velocities)

       !Convert distribution function back to h
       call g_adjust(gnew, phinew, bparnew, direction = to_g_gs2)
    end if
  end subroutine save_restart_dist_fn

  !> Deallocate the [[gs2_diagnostics]] module-level arrays
  subroutine deallocate_arrays
    use mp, only: proc0
    use gs2_heating, only: del_htype
    use diagnostics_fluxes, only: finish_diagnostics_fluxes
    use diagnostics_kinetic_energy_transfer, only: finish_diagnostics_kinetic_energy_transfer
    use diagnostics_nonlinear_convergence, only: finish_nonlinear_convergence
    use collisional_heating, only: finish_collisional
    use diagnostics_zonal_transfer, only: finish_diagnostics_transfer
    implicit none

    if (write_heating) then
       call del_htype (h)
       call del_htype (h_hist)
       call del_htype (hk_hist)
       call del_htype (hk)
    end if

    call finish_diagnostics_fluxes()
    call finish_diagnostics_kinetic_energy_transfer()
    call finish_collisional
    call finish_diagnostics_transfer
    if (allocated(h_hist)) deallocate (h_hist, hk_hist, hk)
    if (allocated(j_ext_hist)) deallocate (j_ext_hist)
    if (proc0 .and. allocated(omegahist)) deallocate (omegahist)

    if (allocated(omega)) deallocate(omega, omegaavg)
    if (use_nonlin_convergence) call finish_nonlinear_convergence

    if (proc0 .and. allocated(domega)) deallocate(domega)
  end subroutine deallocate_arrays

  !> Close various text output files
  subroutine close_files
    use file_utils, only: close_output_file
    use mp, only: proc0
    implicit none

    if(.not.proc0) return
    if (write_ascii .and. write_parity) call close_output_file (parity_unit)
    if (write_ascii .and. write_verr) then
       call close_output_file (res_unit)
       call close_output_file (lpc_unit)
       call close_output_file (res_unit2)
    end if
    if (write_ascii .and. write_cerr) call close_output_file(cres_unit)
    if (write_ascii .and. write_g) call close_output_file(dist_unit)
    if (write_ascii .and. write_gyx) call close_output_file(yxdist_unit)
    if (write_ascii) call close_output_file (out_unit)
    if (write_ascii .and. write_cross_phase) call close_output_file (phase_unit)
    if (write_ascii .and. write_heating) call close_output_file (heat_unit)
    if (write_ascii .and. write_heating) call close_output_file (heat_unit2)
    if (write_ascii .and. write_jext) call close_output_file (jext_unit)
  end subroutine close_files

  !> Calculate and write the various diagnostic quantities.
  !>
  !> This is the main routine for the "old" diagnostics.
  subroutine loop_diagnostics (istep, exit, debopt, force)
    use, intrinsic :: iso_fortran_env, only: output_unit
    use species, only: spec, has_electron_species
    use kt_grids, only: naky, ntheta0
    use run_parameters, only: nstep, has_phi, has_apar, has_bpar
    use fields_arrays, only: phinew, bparnew
    use collisions, only: ncheck, vary_vnew
    use mp, only: proc0, broadcast
    use gs2_time, only: user_time
    use gs2_io, only: nc_final_fields, nc_sync, nc_write_lorentzian
    use antenna, only: antenna_w
    use optionals, only: get_option_with_default
    use diagnostics_printout, only: write_phi_fluxes_to_unit, write_fluxes_to_unit
    use diagnostics_fluxes, only: calculate_fluxes, species_energy_exchange, &
         species_es_heat_flux, species_apar_heat_flux, species_bpar_heat_flux, &
         species_es_particle_flux, species_apar_particle_flux, species_bpar_particle_flux, &
         species_es_momentum_flux, species_heat_flux, species_particle_flux
    use diagnostics_velocity_space, only: collision_error
    use optionals, only: get_option_with_default
    use warning_helpers, only: is_not_zero, not_exactly_equal
    use collisional_heating, only: write_coll_heat => write_collisional, calculate_collisional
    use diagnostics_zonal_transfer, only: write_zonal_trans => write_zonal_transfer, calculate_zonal_transfer
    implicit none
    !> Current timestep
    integer, intent (in) :: istep
    !> If true, simulation should stop (for example, because it has converged)
    logical, intent (out) :: exit
    !> If true, turn on some debug prints
    logical, intent (in), optional:: debopt
    logical, intent(in), optional :: force
    real, dimension (ntheta0, naky) :: phitot
    complex :: wtmp_new
    integer :: ik, it, write_mod
    real :: phi2, apar2, bpar2, t
    logical :: do_force, debug
    real, save :: t_old = 0.
    integer, save :: istep_last = -1
    debug = get_option_with_default(debopt, .false.)

    !Set the current time
    t = user_time
    !Always skip if we've already done this step
    if (istep == istep_last) return

    do_force = get_option_with_default(force, .false.)

    exit = .false.

    call do_get_omega(istep,debug,exit)

    if (write_heating) call heating (istep, navg, h, hk, h_hist, hk_hist)

    if (make_movie .and. mod(istep,nmovie) == 0) call do_write_movie(t)

    if (write_gyx .and. mod(istep,nmovie) == 0) call do_write_fyx (yxdist_unit, phinew, bparnew)

    if (vary_vnew) then
       write_mod = mod(istep,ncheck)
    else
       write_mod = mod(istep,nwrite)
    end if

    if (write_verr .and. write_mod == 0) call do_write_verr

    if (debug) write(6,*) "loop_diagnostics: call update_time"

    !########################################################
    !The rest of the routine only happens every nwrite steps
    !########################################################
    if (mod(istep,nwrite) /= 0 .and. .not. exit .and. .not. do_force) return

    ! If we get here then we're doing a full set of diagnostics
    ! so store the step
    istep_last = istep

    if (write_g) call do_write_f (dist_unit)

    ! Note this also returns phi2, apar2, bpar2 and phitot for other diagnostics
    if (proc0) call do_write_ncloop(t, istep, phi2, apar2, bpar2, phitot)

    if (print_line) call do_print_line(phitot)

    if (write_moments) call do_write_moments

    if (write_cross_phase .and. has_electron_species(spec) .and. write_ascii) call do_write_crossphase(t)

    !###########################
    ! The following large section could do with being moved to separate routines but
    !     at the moment its all quite interlinked which makes this hard.
    !###########################

    if (debug) write(6,*) "loop_diagnostics: -1"

    if (write_any_fluxes) then
       call calculate_fluxes(get_netcdf_file_id(), nout, write_fluxes, write_fluxes_by_mode, &
            write_energy_exchange, write_fluxes_by_energy)
    end if

    if (proc0 .and. print_flux_line) then
       if (has_phi) call write_phi_fluxes_to_unit(output_unit, t, phi2, &
            species_es_heat_flux, species_energy_exchange)
       if (has_apar) call write_fluxes_to_unit(output_unit, t, apar2, "apar", &
            species_apar_heat_flux)
       if (has_bpar) call write_fluxes_to_unit(output_unit, t, bpar2, "bpar", &
            species_bpar_heat_flux)
    end if

    ! Check for convergence
    if (use_nonlin_convergence) call check_nonlin_convergence(istep, &
         species_es_heat_flux(1), exit)
    if (debug) write(6,*) "loop_diagnostics: -1"

    if (.not. write_any) then
       ! We have to increment the number of outputs written to file
       ! before leaving. Usually we do this at the end of the routine
       ! so we must make sure we do this here before leaving early.
       nout = nout + 1
       return
    end if

    if (debug) write(output_unit, *) "loop_diagnostics: -2"

    if (proc0 .and. write_any) then
       if (write_ascii) then
         write (unit=out_unit, fmt=*) 'time=', t
         if (write_heating) call do_write_heating(t, heat_unit, heat_unit2, h)
         if (write_jext) call do_write_jext(t,istep)
       end if

       if (write_flux_line .and. write_ascii) then
          if (has_phi) call write_phi_fluxes_to_unit(out_unit, t, phi2, &
               species_es_heat_flux, species_energy_exchange, species_es_particle_flux, &
               species_es_momentum_flux)

          if (write_lorentzian) then
             wtmp_new = antenna_w()
             call nc_write_lorentzian(get_netcdf_file_id(), nout, wtmp_new)

             if (is_not_zero(real(wtmp_old)) .and. not_exactly_equal(wtmp_new, wtmp_old)) &
                  write (unit=out_unit, fmt="('w= ',e17.10, ' amp= ',e17.10)") real(wtmp_new), sqrt(2.*apar2)
             wtmp_old = wtmp_new
          end if
          
          if (has_apar) call write_fluxes_to_unit(out_unit, t, apar2, "apar", &
               species_apar_heat_flux, species_apar_particle_flux)
          
          if (has_bpar) call write_fluxes_to_unit(out_unit, t, bpar2, "bpar", &
               species_bpar_heat_flux, species_bpar_particle_flux)

          write (unit=out_unit, fmt="('t= ',e17.10,' h_tot= ',e11.4, ' z_tot= ',e11.4)") t, sum(species_heat_flux), sum(species_particle_flux * spec%z)
       end if

       if (write_ascii) then
          do ik = 1, naky
             do it = 1, ntheta0
                if (write_line) call do_write_line(t,it,ik,phitot(it,ik))
                if (write_omega) call do_write_omega(it,ik)
                if (write_omavg) call do_write_omavg(it,ik)
             end do
          enddo
       end if
    endif

    if (write_cerr) call collision_error(cres_unit, phinew, bparnew)

    if (write_nl_flux_dist) call do_write_nl_flux_dist(get_netcdf_file_id(), nout)

    if (write_kinetic_energy_transfer) call do_write_kinetic_energy_transfer(get_netcdf_file_id(), nout)

    if (write_symmetry) call do_write_symmetry(get_netcdf_file_id(), nout)

    if (write_correlation) call do_write_correlation(get_netcdf_file_id(), nout)

    if (write_correlation_extend .and. istep > nstep/4 .and. mod(istep,nwrite_mult*nwrite)==0) &
         call do_write_correlation_extend(get_netcdf_file_id(), t, t_old)

    if (write_parity) call do_write_parity(t, parity_unit, write_ascii)

    if (write_collisional) then
       call calculate_collisional()
       call write_coll_heat(get_netcdf_file_id(), nout)
    end if

    if (write_zonal_transfer) then
       call calculate_zonal_transfer
       call write_zonal_trans(get_netcdf_file_id(), nout)
    end if

    !Now sync the data to file (note doesn't actually sync every call)
    if (proc0) call nc_sync(get_netcdf_file_id(), nout, nc_sync_freq)

    !Increment loop counter
    nout = nout + 1

    if (write_ascii .and. mod(nout, 10) == 0 .and. proc0) call flush_files

    !Update time
    t_old = t

    if (debug) write(6,*) "loop_diagnostics: done"
  end subroutine loop_diagnostics

  !> Calculate [[gs2_diagnostics_knobs:omegaavg]] for linear simulations or if
  !> [[kt_grids:eqzip]] is on; otherwise set
  !> [[gs2_diagnostics_knobs:omega]] and [[gs2_diagnostics_knobs:omegaavg]] to
  !> zero
  subroutine do_get_omega(istep,debug,exit)
    use mp, only: proc0, broadcast
    use nonlinear_terms, only: nonlin
    use kt_grids, only: ieqzip
    implicit none
    !> The current timestep
    integer, intent(in) :: istep
    !> Turn on some debug messages
    logical, intent(in) :: debug
    !> Returns true if the simulation has converged (see
    !> [[gs2_diagnostics_knobs:omegatol]]) or if a numerical instability has
    !> occurred (see [[gs2_diagnostics_knobs:omegainst]]).
    logical, intent(inout) :: exit

    if (nonlin .and. .not. any(ieqzip)) then
       !Make sure we've at least initialised the omega arrays
       !for any later output etc.
       omega = 0.
       omegaavg = 0.
    else
       if (proc0) then
          if (debug) write(6,*) "loop_diagnostics: proc0 call get_omegaavg"
          call get_omegaavg (istep, exit, omegaavg, debug)
          if (debug) write(6,*) "loop_diagnostics: proc0 done called get_omegaavg"
       endif
       call broadcast (exit)
    end if

  end subroutine do_get_omega

  !> Compute volume averages of the fields and write the fields, field averages,
  !> heating and frequency to the netCDF files
  subroutine do_write_ncloop(time, istep, phi2, apar2, bpar2, phitot)
    use gs2_time, only: woutunits, tunits
    use gs2_io, only: nc_loop
    use diagnostics_printout, only: phinorm
    use mp, only: proc0
    use kt_grids, only: ntheta0, naky
    implicit none
    !> Simulation time
    real, intent(in) :: time
    !> Current timestep
    integer, intent(in) :: istep
    !> Fields squared
    real, intent(out) :: phi2, apar2, bpar2
    !> FIXME: add documentation. Needs [[phinorm]] documenting
    real, dimension (:, :), intent(out) :: phitot
    real, dimension (ntheta0, naky) :: ql_metric, growth_rates
    phitot = 0.0
    if (.not. proc0) return

    omega = omegahist(mod(istep,navg),:,:)

    call phinorm (phitot)

    if (write_ql_metric) then
       ! Calculate the instantaneous omega and keep the growth rate.
       growth_rates = aimag(calculate_instantaneous_omega())
       ql_metric = calculate_simple_quasilinear_flux_metric_by_k(growth_rates)
    else
       ql_metric = 0.0
    end if

    call nc_loop (get_netcdf_file_id(), nout, time, phi2, apar2, bpar2, igomega, &
         h, hk, omega, omegaavg, woutunits, tunits, phitot, ql_metric, write_omega, &
         write_heating, write_ql_metric, write_phi_over_time, write_apar_over_time, &
         write_bpar_over_time)
  end subroutine do_write_ncloop

  !> Write \(\phi, A_\parallel, B_\parallel\) normalised to the value of
  !> \(\phi\) at the outboard midplane
  !>
  !> Writes to text file `<runname>.eigenfunc` and/or netCDF
  subroutine do_write_eigenfunc(file_id, write_text)
    use mp, only: proc0
    use file_utils, only: open_output_file, close_output_file
    ! This renaming seems needed to avoid an issue with ifx
    use diagnostics_fields, only: write_eigenfunc_ => write_eigenfunc
    implicit none
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> If true, write text file
    logical, intent(in) :: write_text
    integer :: unit
    if (.not. proc0) return
    if (write_text) call open_output_file (unit, ".eigenfunc")
    call write_eigenfunc_(write_text, unit, file_id)
    if (write_text) close(unit)
  end subroutine do_write_eigenfunc

  !> Write some geometry information to text file `<runname>.g`
  subroutine do_write_geom
    use mp, only: proc0
    use file_utils, only: open_output_file, close_output_file
    use theta_grid, only: ntgrid, theta, Rplot, Zplot, aplot, Rprime, Zprime, aprime, drhodpsi, qval, shape
    implicit none
    integer :: i, unit

    if(.not.proc0) return

    !May want to guard this with if(write_ascii) but not until we provide this
    !data in main netcdf output by default.
    call open_output_file (unit, ".g")
    write (unit,fmt="('# shape: ',a)") trim(shape)
    write (unit,fmt="('# q = ',e11.4,' drhodpsi = ',e11.4)") qval, drhodpsi
    write (unit,fmt="('# theta1             R2                  Z3               alpha4      ', &
         &   '       Rprime5              Zprime6           alpha_prime7 ')")
    do i=-ntgrid,ntgrid
       write (unit,1000) theta(i),Rplot(i),Zplot(i),aplot(i), &
            Rprime(i),Zprime(i),aprime(i)
    enddo
    call close_output_file (unit)
1000 format(20(1x,1pg18.11))
  end subroutine do_write_geom

  !> Transform \(\phi, A_\parallel, B_\parallel\) to real space, then write to netCDF
  subroutine do_write_movie(time)
    use gs2_io, only: nc_loop_movie, get_netcdf_file_id
    use mp, only: proc0
    implicit none
    !> Current simulation time
    real, intent(in) :: time
    if (proc0) call nc_loop_movie(get_netcdf_file_id(), nout_movie, time)
    nout_movie = nout_movie + 1
  end subroutine do_write_movie

  !> Print estimated frequencies and growth rates to screen/stdout
  subroutine do_print_line(phitot)
    use mp, only: proc0
    use kt_grids, only: naky, ntheta0, aky, akx, theta0
    use gs2_time, only: woutunits
    implicit none
    !> Total magnitude of all the fields
    real, dimension (:, :), intent(in) :: phitot
    integer :: ik, it

    if(.not.proc0) return
    do ik = 1, naky
       do it = 1, ntheta0
          write (unit=*, fmt="('ky=',1pe9.2, ' kx=',1pe9.2, &
               & ' om=',e9.2,1x,e9.2,' omav=',e9.2,1x,e9.2, &
               & ' phtot=',e9.2,' theta0=',1pe9.2)") &
               aky(ik), akx(it), &
               real( omega(it,ik)*woutunits(ik)), &
               aimag(omega(it,ik)*woutunits(ik)), &
               real( omegaavg(it,ik)*woutunits(ik)), &
               aimag(omegaavg(it,ik)*woutunits(ik)), &
               phitot(it,ik), theta0(it,ik)
       end do
    end do
    write (*,*) 
  end subroutine do_print_line

  !> Calculate and write the time-averaged antenna current to [[jext_unit]]
  subroutine do_write_jext(time, istep)
    use kt_grids, only: ntheta0, naky
    use mp, only: proc0
    use warning_helpers, only: is_not_zero
    use gs2_io, only: nc_write_jext
    implicit none
    !> Current simulation time
    real, intent(in) :: time
    !> Current timestep
    integer, intent(in) :: istep
    integer :: ik, it
    real, dimension(:,:), allocatable ::  j_ext

    if (.not. proc0) return
    allocate (j_ext(ntheta0, naky)); j_ext=0.
    call calc_jext(istep, j_ext)
    call nc_write_jext(get_netcdf_file_id(), nout, j_ext)
    if (.not. write_ascii) return

    do ik=1,naky
       do it = 1, ntheta0
          if (is_not_zero(j_ext(it, ik))) then
             write (unit=jext_unit, fmt="(es12.4,i4,i4,es12.4)")  &
                  time,it,ik,j_ext(it,ik)
          endif
       enddo
    enddo
    deallocate(j_ext)
  end subroutine do_write_jext

  !> Write various moments to netCDF
  subroutine do_write_moments
    use diagnostics_moments, only: write_moments
    implicit none
    call write_moments(get_netcdf_file_id(), nout, igomega, write_moments_over_time, &
         write_moments_by_mode)
  end subroutine do_write_moments

  !> Calculate the cross-phase (see [[get_cross_phase]]) and write to the
  !> [[phase_unit]] file
  subroutine do_write_crossphase(time)
    use mp, only: proc0
    use diagnostics_turbulence, only: get_cross_phase
    use gs2_io, only: nc_write_cross_phase
    implicit none
    !> Current simulation time
    real, intent(in) :: time
    real :: phase_tot, phase_theta
    call get_cross_phase (phase_tot, phase_theta)
    if (.not. proc0) return
    if (write_ascii) write (unit=phase_unit, &
         fmt="('t= ',e17.10,' phase_tot= ',e11.4,' phase_theta= ',e11.4)") &
         time, phase_tot, phase_theta
    call nc_write_cross_phase(get_netcdf_file_id(), nout, phase_theta, phase_tot)
  end subroutine do_write_crossphase

  !> Write estimated frequency and growth rates to [[out_unit]] for an individual \((k_y, \theta)\) point
  subroutine do_write_line(time,it,ik,phitot)
    use kt_grids, only: aky, akx, theta0
    use gs2_time, only: woutunits
    implicit none
    !> Simulation time
    real, intent(in) :: time, phitot
    !> \(k_y\)- and \(\theta\)-indices
    integer, intent(in) :: ik, it
    if (.not. write_ascii) return
    write (out_unit, "('t= ',e17.10,' aky= ',1p,e12.4, ' akx= ',1p,e12.4, &
         &' om= ',1p,2e12.4,' omav= ', 1p,2e12.4,' phtot= ',1p,e12.4,' theta0= ',1p,e12.4)") &
         time, aky(ik), akx(it), &
         real( omega(it,ik)*woutunits(ik)), &
         aimag(omega(it,ik)*woutunits(ik)), &
         real( omegaavg(it,ik)*woutunits(ik)), &
         aimag(omegaavg(it,ik)*woutunits(ik)), &
         phitot,theta0(it,ik)
  end subroutine do_write_line

  !> Write instantaneous growth rate and frequency to [[out_unit]] for an individual \((k_y, \theta)\) point
  subroutine do_write_omega(it,ik)
    use gs2_time, only: tunits, woutunits
    implicit none
    !> \(k_y\)- and \(\theta\)-indices
    integer, intent(in) :: it, ik
    if (.not. write_ascii) return
    write (out_unit,&
         fmt='(" omega= (",1p,e12.4,",",1p,e12.4,")",t45,"omega/(vt/a)= (",1p,e12.4,",",1p,e12.4,")")') &
         omega(it,ik)/tunits(ik), omega(it,ik)*woutunits(ik)
  end subroutine do_write_omega

  !> Write time-averaged growth rate and frequency to [[out_unit]] for an individual \((k_y, \theta)\) point
  subroutine do_write_omavg(it,ik)
    use gs2_time, only: tunits, woutunits
    implicit none
    integer, intent(in) :: it, ik
    if (.not. write_ascii) return
    write (out_unit,&
         fmt='(" omavg= (",1p,e12.4,",",1p,e12.4,")",t45,"omavg/(vt/a)= (",1p,e12.4,",",1p,e12.4,")")') &
         omegaavg(it,ik)/tunits(ik), omegaavg(it,ik)*woutunits(ik)               
  end subroutine do_write_omavg

  !> Write some error estimates.
  !>
  !> Error estimates are calculated by:
  !> - comparing standard integral with less-accurate integral
  !> - monitoring amplitudes of legendre polynomial coefficients
  !>
  !> Only writes to text file, requires [[lpc_unit]] to be `open`
  subroutine do_write_verr
    use mp, only: proc0
    use le_grids, only: grid_has_trapped_particles
    use fields_arrays, only: phinew, bparnew
    use gs2_time, only: user_time
    use collisions, only: vnmult
    use species, only: spec
    use diagnostics_velocity_space, only: get_verr, get_gtran
    use gs2_io, only: nc_write_verr_est
    implicit none
    real, dimension(5,2) :: errest
    integer, dimension (5,3) :: erridx
    real :: geavg, glavg, gtavg

    errest = 0.0; erridx = 0

    ! error estimate obtained by comparing standard integral with less-accurate integral
    call get_verr (errest, erridx, phinew, bparnew)

    ! error estimate based on monitoring amplitudes of legendre polynomial coefficients
    call get_gtran (geavg, glavg, gtavg, phinew, bparnew)

    if (proc0) then
       call nc_write_verr_est(get_netcdf_file_id(), nout, &
         [geavg, glavg, gtavg], errest, [vnmult(1)*spec(1)%vnewk, vnmult(2)*spec(1)%vnewk], &
         erridx)
       ! write error estimates for ion dist. fn. at outboard midplane with ik=it=1 to ascii files
       if (write_ascii) then
          if (grid_has_trapped_particles()) then
             write(lpc_unit,"(4(1x,e13.6))") user_time, geavg, glavg, gtavg
          else
             write(lpc_unit,"(3(1x,e13.6))") user_time, geavg, glavg
          end if
          write(res_unit,"(8(1x,e13.6))") user_time, errest(1,2), errest(2,2), errest(3,2), &
               errest(4,2), errest(5,2), vnmult(1)*spec(1)%vnewk, vnmult(2)*spec(1)%vnewk
          !Max verr
          write(res_unit2,"(3(i8),(1x,e13.6),3(i8),(1x,e13.6),3(i8),(1x,e13.6),3(i8),(1x,e13.6),3(i8),(1x,e13.6))") &
               erridx(1,1), erridx(1,2), erridx(1,3), errest(1,1), &
               erridx(2,1), erridx(2,2), erridx(2,3), errest(2,1), &
               erridx(3,1), erridx(3,2), erridx(3,3), errest(3,1), &
               erridx(4,1), erridx(4,2), erridx(4,3), errest(4,1), &
               erridx(5,1), erridx(5,2), erridx(5,3), errest(5,1)
       end if
    end if
  end subroutine do_write_verr

  !> Write the (total) heating diagnostics to [[heat_unit]] and [[heat_unit2]] (per-species)
  subroutine do_write_heating(t, file_unit, file_unit2, h)
    use species, only: nspec
    implicit none
    real, intent(in) :: t
    integer, intent(in) :: file_unit, file_unit2
    !> Heating diagnostics
    type(heating_diagnostics), intent (in) :: h

    integer :: is

    !
    ! For case with two species:
    !
    ! Column     Item               
    !   1        time              
    !   2        Energy              
    !   3        dEnergy/dt            
    !   4        J_ant.E             
    !   5        [h_(i+1)*h_*]/2 * C[h_(i+1)] * T_0 for species 1
    !   6        [h_(i+1)*h_*]/2 * C[h_(i+1)] * T_0 for species 2
    !   7       -[h H(h) * T_0]_1
    !   8       -[h H(h) * T_0]_2
    !   9       -[h C(h) * T_0]_1 
    !  10       -[h C(h) * T_0]_2
    !  11        [h w_* h]_1
    !  12        [h w_* h]_2
    !  13        [h * (q dchi/dt - dh/dt * T0)]_1
    !  14        [h * (q dchi/dt - dh/dt * T0)]_2
    !  15      sum (h C(h) * T_0)  in total, as in 5, 6      
    !  16     -sum (h H(h) * T_0)      
    !  17     -sum (h C(h) * T_0)   
    !  18      sum (h w_* h)  
    !  19      sum [h (q dchi/dt - dh/dt * T0)]
    !  20      3 + 4 + 18 + 19
    !  21      (k_perp A)**2
    !  22      B_par**2
    !  23      df_1 ** 2
    !  24      df_2 ** 2
    !  25      h_1 ** 2
    !  26      h_2 ** 2
    !  27      Phi_bar_1 ** 2
    !  28      Phi_bar_2 ** 2
    !
    !
    ! For case with one species:
    !
    ! Column     Item               
    !   1        time              
    !   2        Energy              
    !   3        dEnergy/dt            
    !   4        J_ant.E             
    !   5        [h_(i+1)*h_*]/2 * C[h_(i+1)] * T_0 
    !   6       -[h H(h) * T_0]
    !   7       -[h C(h) * T_0]
    !   8        [h w_* h]
    !   9        [h * (q dchi/dt - dh/dt * T0)]_1
    !  10      sum (h C(h) * T_0)  in total, as in 5, 6      
    !  11     -sum (h H(h) * T_0)      
    !  12     -sum (h C(h) * T_0)   
    !  13      sum (h w_* h)  
    !  14      sum [h (q dchi/dt - dh/dt * T0)]
    !  15      3 + 4 + 9 + 10
    !  16      (k_perp A)**2
    !  17      B_par**2
    !  18      df ** 2
    !  19      h ** 2
    !  20      Phi_bar ** 2
    
    write (unit=file_unit, fmt="(28es12.4)") t,h % energy,  &
         h % energy_dot, h % antenna, h % imp_colls, h % hypercoll, h % collisions, &
         h % gradients, h % heating, sum(h % imp_colls), sum(h % hypercoll), sum(h % collisions), &
         sum(h % gradients), sum(h % heating),sum(h%heating)+h%antenna+sum(h%gradients)+h%energy_dot, &
         h % eapar, h % ebpar, h % delfs2(:),  h % hs2(:), h % phis2(:)

    do is=1,nspec
      write (unit=file_unit2, fmt="(15es12.4)") t,h % energy,  &
           h % energy_dot, h % antenna, h % imp_colls(is), h % hypercoll(is), h % collisions(is), &
           h % gradients(is), h % heating(is), &
           h % eapar, h % ebpar, h % delfs2(is),  h % hs2(is), h % phis2(is), real(is)
      write (unit=file_unit2, fmt=*)
    end do
    write (unit=file_unit2, fmt=*)

    flush (file_unit)
    flush (file_unit2)
  end subroutine do_write_heating

  !> Calculate the momentum flux as a function of \((v_\parallel, \theta, t)\)
  !> and write to netCDF
  subroutine do_write_symmetry(file_id, nout)
    use diagnostics_fluxes, only: flux_vs_theta_vs_vpa
    use theta_grid, only: ntgrid
    use le_grids, only: nlambda, negrid
    use species, only: nspec
    use mp, only: proc0
    use gs2_io, only: nc_loop_sym
    use fields_arrays, only: phinew
    implicit none
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> Current timestep
    integer, intent(in) :: nout

    real, dimension(:,:,:), allocatable :: vflx_sym
    allocate (vflx_sym(-ntgrid:ntgrid,nlambda*negrid,nspec))
    call flux_vs_theta_vs_vpa (phinew, vflx_sym)
    if (proc0) call nc_loop_sym (file_id, nout, vflx_sym)
    deallocate (vflx_sym)
  end subroutine do_write_symmetry

  subroutine do_write_kinetic_energy_transfer(file_id, nout)
    use diagnostics_kinetic_energy_transfer, only: calculate_kinetic_energy_transfer, write_kinetic_energy_transfer
    use mp, only: proc0
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> Current timestep
    integer, intent(in) :: nout
    call calculate_kinetic_energy_transfer
    if (proc0) call write_kinetic_energy_transfer(file_id, nout)
  end subroutine do_write_kinetic_energy_transfer

  !> Calculate the poloidal distributions of the fluxes of particles, parallel
  !> momentum, perpendicular momentum, and energy (see section 3.1 and appendix
  !> A of [Ball et al. PPCF 58 (2016)
  !> 045023](https://doi.org/10.1088/0741-3335/58/4/045023) as well as section 5
  !> of "GS2 analytic geometry specification")
  subroutine do_write_nl_flux_dist(file_id, nout)
    use diagnostics_fluxes, only: flux_dist
    use theta_grid, only: ntgrid
    use species, only: nspec, spec
    use mp, only: proc0
    use gs2_io, only: nc_loop_dist
    use kt_grids, only: naky, ntheta0
    use fields_arrays, only: phinew, bparnew
    use run_parameters, only: has_phi
    implicit none
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> Current timestep
    integer, intent(in) :: nout

    real, dimension (:,:), allocatable :: part_fluxes_dist, mom_fluxes_par_dist
    real, dimension (:,:), allocatable :: phi_dist, mom_fluxes_perp_dist, heat_fluxes_dist
    real, dimension (:,:,:,:), allocatable :: pflux_dist, vflux_par_dist, vflux_perp_dist, qflux_dist
    integer :: ig, is
    if (.not. has_phi) return

    allocate (pflux_dist(-ntgrid:ntgrid,ntheta0,naky,nspec))
    allocate (vflux_par_dist(-ntgrid:ntgrid,ntheta0,naky,nspec))
    allocate (vflux_perp_dist(-ntgrid:ntgrid,ntheta0,naky,nspec))
    allocate (qflux_dist(-ntgrid:ntgrid,ntheta0,naky,nspec))

    call flux_dist (phinew, bparnew, pflux_dist, vflux_par_dist, vflux_perp_dist, qflux_dist)

    if (proc0) then
       allocate (part_fluxes_dist(-ntgrid:ntgrid,nspec))
       allocate (mom_fluxes_par_dist(-ntgrid:ntgrid,nspec))
       allocate (mom_fluxes_perp_dist(-ntgrid:ntgrid,nspec))
       allocate (heat_fluxes_dist(-ntgrid:ntgrid,nspec))
       allocate (phi_dist(2,-ntgrid:ntgrid))
       do is = 1, nspec
          do ig = -ntgrid, ntgrid
             call get_volume_average (pflux_dist(ig,:,:,is), part_fluxes_dist(ig,is))
             call get_volume_average (vflux_par_dist(ig,:,:,is), mom_fluxes_par_dist(ig,is))
             call get_volume_average (vflux_perp_dist(ig,:,:,is), mom_fluxes_perp_dist(ig,is))
             call get_volume_average (qflux_dist(ig,:,:,is), heat_fluxes_dist(ig,is))
          end do

          part_fluxes_dist(:, is) = part_fluxes_dist(:, is) * spec(is)%dens
          mom_fluxes_par_dist(:, is) = mom_fluxes_par_dist(:, is) * spec(is)%dens * sqrt(spec(is)%mass * spec(is)%temp)             
          mom_fluxes_perp_dist(:, is) = mom_fluxes_perp_dist(:, is) * spec(is)%dens * sqrt(spec(is)%mass * spec(is)%temp)
          heat_fluxes_dist(:, is) = heat_fluxes_dist(:, is) * spec(is)%dens * spec(is)%temp
       end do

       do ig = -ntgrid, ntgrid
          call get_volume_average (real(phinew(ig,:,:)), phi_dist(1,ig))
          call get_volume_average (aimag(phinew(ig,:,:)), phi_dist(2,ig))
       end do

       call nc_loop_dist (file_id, nout, part_fluxes_dist, mom_fluxes_par_dist, &
            mom_fluxes_perp_dist, heat_fluxes_dist, phi_dist)

       deallocate (part_fluxes_dist, mom_fluxes_par_dist, mom_fluxes_perp_dist)
       deallocate (heat_fluxes_dist, phi_dist)
     end if
    deallocate (pflux_dist, vflux_par_dist, vflux_perp_dist, qflux_dist)
   end subroutine do_write_nl_flux_dist

  !> Calculate the correlation function over the physical domain and write to netCDF
  subroutine do_write_correlation(file_id, nout)
    use theta_grid, only: ntgrid
    use kt_grids, only: naky
    use mp, only: proc0
    use gs2_io, only: nc_loop_corr
    implicit none
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> Current timestep
    integer, intent(in) :: nout
    complex, dimension(:,:), allocatable :: phi_corr_2pi
    allocate (phi_corr_2pi(-ntgrid:ntgrid,naky))
    call correlation (phi_corr_2pi)
    if (proc0) call nc_loop_corr (file_id, nout, phi_corr_2pi)
    deallocate (phi_corr_2pi)
  end subroutine do_write_correlation

  !> Calculate the correlation function over the extended domain and write to netCDF
  !>
  !> This does the calculation every time it is called, but only writes every
  !> `nwrite_mult * nwrite` timesteps
  subroutine do_write_correlation_extend(file_id, time, time_old)
    use theta_grid, only: ntgrid
    use kt_grids, only: jtwist, ntheta0, naky
    use mp, only: proc0
    use gs2_io, only: nc_loop_corr_extend
    implicit none
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> Current and previous simulation times
    real, intent(in) :: time, time_old
    complex, dimension (:,:,:), allocatable, save:: phicorr_sum
    real, dimension (:,:,:), allocatable, save :: phiextend_sum
    complex, dimension (:,:,:), allocatable :: phi_corr
    real, dimension (:,:,:), allocatable :: phi2_extend
    real, save :: tcorr0 = 0.0
    integer :: ntg_extend
    if (.not. proc0) return
    ntg_extend = (2*ntgrid+1) * ((ntheta0-1) / jtwist + 1)

    if (.not. allocated(phicorr_sum)) then
       ! Warning: For ntheta=ntheta0=naky=128 and jtwist = 6 these two arrays
       ! will require 750 MB and 375 MB respectively and they are not freed
       ! after leaving the routine.
       allocate (phicorr_sum(ntg_extend,ntheta0,naky)) ; phicorr_sum = 0.0
       allocate (phiextend_sum(ntg_extend,ntheta0,naky)) ; phiextend_sum = 0.0
       tcorr0 = time_old
    end if

    allocate (phi_corr(ntg_extend,ntheta0,naky))
    allocate (phi2_extend(ntg_extend,ntheta0,naky))
    call correlation_extend (phi_corr, phi2_extend, ntg_extend)
    phicorr_sum = phicorr_sum + phi_corr*(time-time_old)
    phiextend_sum = phiextend_sum + phi2_extend*(time-time_old)
    call nc_loop_corr_extend (file_id, nout_big, time, phicorr_sum/(time-tcorr0), &
         phiextend_sum/(time-tcorr0))
    nout_big = nout_big + 1
    deallocate (phi_corr, phi2_extend)
  end subroutine do_write_correlation_extend

  !> Calculate average parity of distribution function under \((\theta, k_x,
  !> v_\parallel) \to (-\theta, -k_x, -v_\parallel)\), and write to [[parity_unit]]
  subroutine do_write_parity(t, file_unit, write_text)
    use theta_grid, only: ntgrid, field_line_average
    use kt_grids, only: naky, ntheta0
    use le_grids, only: negrid, nlambda, integrate_moment, integrate_kysum
    use species, only: nspec
    use gs2_layouts, only: init_parity_layouts, p_lo, g_lo, idx_local
    use gs2_layouts, only: ik_idx, it_idx, il_idx, ie_idx, is_idx, idx, proc_id
    use mp, only: send, receive, proc0, iproc, nproc
    use dist_fn_arrays, only: gnew, g_adjust, aj0, to_g_gs2, from_g_gs2
    use fields_arrays, only: phinew, bparnew
    use constants, only: zi
    implicit none
    real, intent(in) :: t
    !> Open formatted file to write text to
    integer, intent(in) :: file_unit
    !> If false, don't calculate or write parity
    logical, intent(in) :: write_text
    integer :: iplo, iglo, sgn2, isgn, il, ie, ig, ik, it, is, it_source
    complex, dimension (:,:,:,:), allocatable :: gparity, gmx, gpx
    complex, dimension (:,:,:), allocatable :: g0, gm, gp
    complex, dimension (:,:,:), allocatable :: g_kint, gm_kint, gp_kint
    complex, dimension (:,:), allocatable :: g_avg, gnorm_avg, phim
    complex, dimension (:), allocatable :: g_all_tot, g_nokx_tot, g_nosig_tot, gtmp
    complex, dimension (:), allocatable :: gnorm_all_tot, gnorm_nokx_tot, gnorm_nosig_tot
    real, dimension (:,:,:), allocatable :: gmnorm, gmint, gpint, gmnormint, gmavg, gpavg, gmnormavg
    real, dimension (:), allocatable :: gmtot, gptot, gtot, gmnormtot
    real, dimension (:), allocatable :: gm_nokx_tot, gp_nokx_tot, gmnorm_nokx_tot
    real, dimension (:), allocatable :: gm_nosig_tot, gp_nosig_tot, gmnorm_nosig_tot

    if (.not. write_text) return

    ! initialize layouts for parity diagnostic. Only does work on the first call
    call init_parity_layouts (naky, nlambda, negrid, nspec, nproc, iproc)

    allocate (gparity(-ntgrid:ntgrid,ntheta0,2,p_lo%llim_proc:p_lo%ulim_alloc))
    allocate (g0(-ntgrid:ntgrid,2,p_lo%llim_proc:p_lo%ulim_alloc))
    allocate (gm(-ntgrid:ntgrid,2,p_lo%llim_proc:p_lo%ulim_alloc))
    allocate (gp(-ntgrid:ntgrid,2,p_lo%llim_proc:p_lo%ulim_alloc))
    allocate (gmnorm(-ntgrid:ntgrid,2,p_lo%llim_proc:p_lo%ulim_alloc))
    allocate (gmint(-ntgrid:ntgrid,naky,nspec))
    allocate (gpint(-ntgrid:ntgrid,naky,nspec))
    allocate (gmnormint(-ntgrid:ntgrid,naky,nspec))
    allocate (gmavg(ntheta0,naky,nspec))
    allocate (gpavg(ntheta0,naky,nspec))
    allocate (gmnormavg(ntheta0,naky,nspec))
    allocate (gmtot(nspec), gm_nokx_tot(nspec), gm_nosig_tot(nspec))
    allocate (gptot(nspec), gp_nokx_tot(nspec), gp_nosig_tot(nspec))
    allocate (gtot(nspec), gmnormtot(nspec), gmnorm_nokx_tot(nspec), gmnorm_nosig_tot(nspec))
    allocate (phim(-ntgrid:ntgrid,naky))
    allocate (g_kint(-ntgrid:ntgrid,2,nspec), gm_kint(-ntgrid:ntgrid,2,nspec), gp_kint(-ntgrid:ntgrid,2,nspec))
    allocate (g_avg(ntheta0,nspec), gnorm_avg(ntheta0,nspec))
    allocate (g_all_tot(nspec), g_nokx_tot(nspec), g_nosig_tot(nspec), gtmp(nspec))
    allocate (gnorm_all_tot(nspec), gnorm_nokx_tot(nspec), gnorm_nosig_tot(nspec))

    ! convert from g to h
    call g_adjust (gnew, phinew, bparnew, direction = from_g_gs2)

    ! below we define gparity = J0 * h, where delta f = h - (q phi / T) F0
    ! because we're ultimately interested in int J0 h * phi (i.e. particle flux)
    ! This should probably be written as a redistribute gather/scatter?
    do iglo = g_lo%llim_world, g_lo%ulim_world
       ik = ik_idx(g_lo, iglo)
       ie = ie_idx(g_lo, iglo)
       il = il_idx(g_lo, iglo)
       is = is_idx(g_lo, iglo)
       it = it_idx(g_lo, iglo)
       ! count total index for given (ik,il,ie,is) combo (dependent on layout)
       iplo = idx(p_lo, ik, il, ie, is)
       ! check to see if value of g corresponding to iglo is on this processor
       if (idx_local(g_lo, iglo)) then
          if (idx_local(p_lo, iplo)) then
             ! if g value corresponding to iplo should be on this processor, then get it
             gparity(:, it, :, iplo) = gnew(:, :, iglo) * spread(aj0(:, iglo), 2, 2)
          else
             ! otherwise, send g for this iplo value to correct processor
             call send (gnew(:, :, iglo) * spread(aj0(:, iglo), 2, 2), proc_id(p_lo, iplo))
          end if
       else if (idx_local(p_lo, iplo)) then
          ! if g for this iplo not on processor, receive it
          call receive (gparity(:, it, :, iplo), proc_id(g_lo, iglo))
       end if
    end do

    ! convert from h back to g
    call g_adjust (gnew, phinew, bparnew, direction = to_g_gs2)

    ! first diagnostic is phase space average of 
    ! |J0*(h(z,vpa,kx) +/- h(-z,-vpa,-kx))|**2 / ( |J0*h(z,vpa,kx)|**2 + |J0*h(-z,-vpa,-kx)|**2 )
    do it = 1, ntheta0
       ! have to treat kx=0 specially
       if (it == 1) then
          it_source = 1
       else
          it_source = ntheta0 - it + 2
       end if

       do isgn = 1, 2
          sgn2 = mod(isgn,2)+1
          do ig = -ntgrid, ntgrid
             ! g0 = gparity(ntgrid:-ntgrid:-1, it_source, 2:1:-1, :) -- copy reversing theta and sigma
             g0(ig, isgn, :) = gparity(-ig, it_source, sgn2, :)
          end do
       end do

       gm = gparity(:,it,:,:)-g0
       gp = gparity(:,it,:,:)+g0
       gmnorm = real(g0*conjg(g0))
       ! integrate out velocity dependence
       call integrate_moment (real(gm*conjg(gm)),gmint,1)
       call integrate_moment (real(gp*conjg(gp)),gpint,1)
       call integrate_moment (gmnorm,gmnormint,1)
       ! average along field line
       gmavg(it,:,:) = field_line_average(gmint)
       gpavg(it,:,:) = field_line_average(gpint)
       gmnormavg(it,:,:) = field_line_average(gmnormint)

       ! phim(theta,kx) = phi(-theta,-kx)
       do ig = -ntgrid, ntgrid
          phim(ig,:) = phinew(-ig, it_source, :)
       end do

       ! want int dtheta sum_{kx} int d3v | sum_{ky} [ J0*(h+ * conjg(phi-) + h- * conjg(phi+)) ky ] |**2
       ! J0*(h+ * conjg(phi-) + h- * conjg(phi+)) = h*conjg(phi) - h(-theta,-vpa,-kx)*conjg(phi(-theta,-kx))
       do iplo = p_lo%llim_proc, p_lo%ulim_proc
          ik = ik_idx(p_lo,iplo)
          do isgn = 1, 2
             gm(:,isgn,iplo) = gparity(:,it,isgn,iplo)*conjg(phinew(:,it,ik)) - g0(:,isgn,iplo)*conjg(phim(:,ik))
          end do
       end do
       call integrate_kysum (gm, g_kint, 1)
       g_avg(it, :) = field_line_average(g_kint(:,1,:) + g_kint(:,2,:))

       ! get normalizing term for above diagnostic
       do iplo = p_lo%llim_proc, p_lo%ulim_proc
          ik = ik_idx(p_lo,iplo)
          do isgn = 1, 2
             gm(:,isgn,iplo) = gparity(:,it,isgn,iplo)*conjg(phinew(:,it,ik))
             gp(:,isgn,iplo) = g0(:,isgn,iplo)*conjg(phim(:,ik))
          end do
       end do
       call integrate_kysum (gm, gm_kint, 1)
       call integrate_kysum (gp, gp_kint, 1)
       gnorm_avg(it,:) = field_line_average(gm_kint(:,1,:)+gm_kint(:,2,:) &
            + gp_kint(:,1,:)+gp_kint(:,2,:))
    end do

    ! average over perp plane
    do is = 1, nspec
       call get_volume_average (gmavg(:,:,is), gmtot(is))
       call get_volume_average (gpavg(:,:,is), gptot(is))
       call get_volume_average (gmnormavg(:,:,is), gmnormtot(is))    
       g_all_tot(is) = sum(g_avg(:,is))
       gnorm_all_tot(is) = sum(gnorm_avg(:,is))
    end do

    allocate (gmx(-ntgrid:ntgrid,ntheta0,2,p_lo%llim_proc:p_lo%ulim_alloc))
    allocate (gpx(-ntgrid:ntgrid,ntheta0,2,p_lo%llim_proc:p_lo%ulim_alloc))

    ! now we want diagnostic of phase space average of
    ! |J0*(h(z,vpa) +/- h(-z,-vpa))|**2 / ( |J0*h(z,vpa)|**2 + |J0*h(-z,-vpa)|**2 )
    do it = 1, ntheta0
       do isgn = 1, 2
          sgn2 = mod(isgn,2)+1
          do ig = -ntgrid, ntgrid
             g0(ig,isgn,:) = gparity(-ig,it,sgn2,:)
          end do
       end do
       gm = gparity(:,it,:,:)-g0
       gp = gparity(:,it,:,:)+g0
       gmnorm = real(g0*conjg(g0))
       ! integrate out velocity dependence
       call integrate_moment (real(gm*conjg(gm)),gmint,1)
       call integrate_moment (real(gp*conjg(gp)),gpint,1)
       call integrate_moment (gmnorm,gmnormint,1)
       ! average along field line
       gmavg(it,:,:) = field_line_average(gmint)
       gpavg(it,:,:) = field_line_average(gpint)
       gmnormavg(it,:,:) = field_line_average(gmnormint)

       ! phim(theta) = phi(-theta)
       ! have to treat kx=0 specially
       do ig = -ntgrid, ntgrid
          phim(ig,:) = phinew(-ig,it,:)
       end do

       ! want int dtheta int d3v | sum_{kx} sum_{ky} [ J0*(h+ * conjg(phi-) + h- * conjg(phi+)) ky ] |**2
       ! J0*(h+ * conjg(phi-) + h- * conjg(phi+)) = h*conjg(phi) - h(-theta,-vpa)*conjg(phi(-theta))
       do iplo = p_lo%llim_proc, p_lo%ulim_proc
          ik = ik_idx(p_lo,iplo)
          do isgn = 1, 2
             gmx(:,it,isgn,iplo) = g0(:,isgn,iplo)*conjg(phim(:,ik))
             gpx(:,it,isgn,iplo) = gparity(:,it,isgn,iplo)*conjg(phinew(:,it,ik)) - gmx(:,it,isgn,iplo) 
          end do
       end do
    end do

    ! sum over kx
    gp = sum(gpx,2)
    deallocate (gpx)
    call integrate_kysum(gp, g_kint, 1)
    g_nokx_tot = field_line_average(g_kint(:,1,:) + g_kint(:,2,:))

    ! get normalizing terms for above diagnostic
    gm = sum(gmx,2)
    deallocate (gmx)
    call integrate_kysum(gm, gm_kint, 1)
    call integrate_kysum(gm + gp, gp_kint, 1)
    gnorm_nokx_tot = field_line_average(gm_kint(:,1,:)+gm_kint(:,2,:) &
         + gp_kint(:,1,:)+gp_kint(:,2,:))

    ! average over perp plane
    do is = 1, nspec
       call get_volume_average (gmavg(:,:,is), gm_nokx_tot(is))
       call get_volume_average (gpavg(:,:,is), gp_nokx_tot(is))
       call get_volume_average (gmnormavg(:,:,is), gmnorm_nokx_tot(is))    
    end do

    ! final diagnostic is phase space average of 
    ! |J0*(h(z,kx) +/- h(-z,-kx))|**2 / ( |J0*h(z,kx)|**2 + |J0*h(-z,-kx)|**2 )
    do it = 1, ntheta0
       ! have to treat kx=0 specially
       if (it == 1) then
          it_source = 1
       else
          it_source = ntheta0 - it + 2
       end if

       do ig = -ntgrid, ntgrid
          g0(ig, :, :) = gparity(-ig, it_source, :, :)
       end do

       gm = gparity(:,it,:,:)-g0
       gp = gparity(:,it,:,:)+g0
       gmnorm = real(g0*conjg(g0))
       ! integrate out velocity dependence
       call integrate_moment (real(gm*conjg(gm)),gmint,1)
       call integrate_moment (real(gp*conjg(gp)),gpint,1)
       call integrate_moment (gmnorm,gmnormint,1)
       ! average along field line
       gmavg(it,:,:) = field_line_average(gmint)
       gpavg(it,:,:) = field_line_average(gpint)
       gmnormavg(it,:,:) = field_line_average(gmnormint)

       ! phim(theta,kx) = phi(-theta,-kx)
       do ig = -ntgrid, ntgrid
          phim(ig, :) = phinew(-ig, it_source, :)
       end do

       ! want int dtheta sum_{kx} int dE dmu | sum_{sigma} sum_{ky} [ J0*(h+ * conjg(phi-) + h- * conjg(phi+)) ky ] |**2
       ! J0*(h+ * conjg(phi-) + h- * conjg(phi+)) = h*conjg(phi) - h(-theta,-kx)*conjg(phi(-theta,-kx))
       do iplo = p_lo%llim_proc, p_lo%ulim_proc
          ik = ik_idx(p_lo,iplo)
          do isgn = 1, 2
             gm(:,isgn,iplo) = gparity(:,it,isgn,iplo)*conjg(phinew(:,it,ik)) - g0(:,isgn,iplo)*conjg(phim(:,ik))
          end do
       end do

       ! Only sets isgn = 1 in the g_kint array here
       call integrate_kysum (spread(gm(:, 1, :) + gm(:, 2, :), 2, 1), g_kint, 1)
       g_avg(it, :) = field_line_average(g_kint(:,1,:))

       ! get normalizing term for above diagnostic
       do iplo = p_lo%llim_proc, p_lo%ulim_proc
          ik = ik_idx(p_lo,iplo)
          do isgn = 1, 2
             gm(:,isgn,iplo) = gparity(:,it,isgn,iplo)*conjg(phinew(:,it,ik))
             gp(:,isgn,iplo) = g0(:,isgn,iplo)*conjg(phim(:,ik))
          end do
       end do

       ! Only sets isgn = 1 in the g*_kint arrays here
       call integrate_kysum (spread(gm(:, 1, :) + gm(:, 2, :), 2, 1), gm_kint, 1)
       call integrate_kysum (spread(gp(:, 1, :) + gp(:, 2, :), 2, 1), gp_kint, 1)
       gnorm_avg(it, :) = field_line_average(gm_kint(:,1,:) + gp_kint(:,1,:))
    end do

    ! average over perp plane
    do is = 1, nspec
       call get_volume_average (gmavg(:,:,is), gm_nosig_tot(is))
       call get_volume_average (gpavg(:,:,is), gp_nosig_tot(is))
       call get_volume_average (gmnormavg(:,:,is), gmnorm_nosig_tot(is))    
       g_nosig_tot(is) = sum(g_avg(:,is))
       gnorm_nosig_tot(is) = sum(gnorm_avg(:,is))
    end do

    deallocate (gparity) ; allocate (gparity(-ntgrid:ntgrid,ntheta0,naky,nspec))
    ! obtain normalization factor = int over phase space of |g|**2
    call g_adjust (gnew, phinew, bparnew, direction = from_g_gs2)
    !Do all processors need to know the full result here? Only proc0 seems to do any writing.
    !If not then remove the last two arguments in the following call.
    call integrate_moment (spread(aj0,2,2)*spread(aj0,2,2)*gnew*conjg(gnew), gparity, .true., full_arr=.true.)
    call g_adjust (gnew, phinew, bparnew, direction = to_g_gs2)
    do is = 1, nspec
       gmavg(:, :, is) = field_line_average(real(gparity(:,:,:,is)))
       call get_volume_average (gmavg(:,:,is), gtot(is))
    end do

    ! normalize g(theta,vpa,kx) - g(-theta,-vpa,-kx) terms
    where (gtot+gmnormtot > epsilon(0.0))
       gmtot = sqrt(gmtot/(gtot+gmnormtot)) ; gptot = sqrt(gptot/(gtot+gmnormtot))
    elsewhere
       gmtot = sqrt(gmtot) ; gptot = sqrt(gptot)
    end where

    where (real(gnorm_all_tot) > epsilon(0.0))
       gtmp = sqrt(real(g_all_tot)/real(gnorm_all_tot))
    elsewhere
       gtmp = sqrt(real(g_all_tot))
    end where
    where (aimag(gnorm_all_tot) > epsilon(0.0))
       g_all_tot = gtmp + zi*sqrt(aimag(g_all_tot)/aimag(gnorm_all_tot))
    elsewhere
       g_all_tot = gtmp + zi*sqrt(aimag(g_all_tot))
    end where

    ! normalize g(theta,vpa) +/- g(-theta,-vpa) terms
    where (gtot+gmnorm_nokx_tot > epsilon(0.0))
       gm_nokx_tot = sqrt(gm_nokx_tot/(gtot+gmnorm_nokx_tot)) ; gp_nokx_tot = sqrt(gp_nokx_tot/(gtot+gmnorm_nokx_tot))
    elsewhere
       gm_nokx_tot = sqrt(gm_nokx_tot) ; gp_nokx_tot = sqrt(gp_nokx_tot)
    end where

    where (real(gnorm_nokx_tot) > epsilon(0.0))
       gtmp = sqrt(real(g_nokx_tot)/real(gnorm_nokx_tot))
    elsewhere
       gtmp = sqrt(real(g_nokx_tot))
    end where
    where (aimag(gnorm_nokx_tot) > epsilon(0.0))
       g_nokx_tot = gtmp + zi*sqrt(aimag(g_nokx_tot)/aimag(gnorm_nokx_tot))
    elsewhere
       g_nokx_tot = gtmp + zi*sqrt(aimag(g_nokx_tot))
    end where

    ! normalize g(theta,kx) +/ g(-theta,-kx) terms
    where (gtot+gmnorm_nosig_tot > epsilon(0.0))
       gm_nosig_tot = sqrt(gm_nosig_tot/(gtot+gmnorm_nosig_tot)) ; gp_nosig_tot = sqrt(gp_nosig_tot/(gtot+gmnorm_nosig_tot))
    elsewhere
       gm_nosig_tot = sqrt(gm_nosig_tot) ; gp_nosig_tot = sqrt(gp_nosig_tot)
    end where

    where (real(gnorm_nosig_tot) > epsilon(0.0))
       gtmp = sqrt(real(g_nosig_tot)/real(gnorm_nosig_tot))
    elsewhere
       gtmp = sqrt(real(g_nosig_tot))
    end where
    where (aimag(gnorm_nosig_tot) > epsilon(0.0))
       g_nosig_tot = gtmp + zi*sqrt(aimag(g_nosig_tot)/aimag(gnorm_nosig_tot))
    elsewhere
       g_nosig_tot = gtmp + zi*sqrt(aimag(g_nosig_tot))
    end where

    if (proc0) write (file_unit,"(19(1x,e12.5))") t, gmtot, gptot, real(g_all_tot), aimag(g_all_tot), &
         real(gnorm_all_tot), aimag(gnorm_all_tot), gm_nokx_tot, gp_nokx_tot, real(g_nokx_tot), aimag(g_nokx_tot), &
         real(gnorm_nokx_tot), aimag(gnorm_nokx_tot), gm_nosig_tot, gp_nosig_tot, &
         real(g_nosig_tot), aimag(g_nosig_tot), real(gnorm_nosig_tot), aimag(gnorm_nosig_tot)

    deallocate (gparity, g0)
    deallocate (gm, gp, gmnorm, gmint, gpint, gmnormint)
    deallocate (gmavg, gpavg, gmnormavg)
    deallocate (gmtot, gm_nokx_tot, gm_nosig_tot)
    deallocate (gptot, gp_nokx_tot, gp_nosig_tot)
    deallocate (gtot, gmnormtot, gmnorm_nokx_tot, gmnorm_nosig_tot)
    deallocate (g_avg, gnorm_avg)
    deallocate (g_kint, gm_kint, gp_kint)
    deallocate (g_all_tot, g_nokx_tot, g_nosig_tot, gtmp)
    deallocate (gnorm_all_tot, gnorm_nokx_tot, gnorm_nosig_tot)
    deallocate (phim)
  end subroutine do_write_parity

  !> Flush text files (only [[out_unit]], [[res_unit]], [[lpc_unit]], and
  !> [[parity_unit]])
  subroutine flush_files
    implicit none
    if (.not. write_ascii) return
    flush (out_unit)
    if (write_verr) then
       flush (res_unit)
       flush (lpc_unit)
    end if
    if (write_cerr) flush (cres_unit)
    if (write_parity) flush (parity_unit)
    if (write_g) flush (dist_unit)
    if (write_gyx) flush (yxdist_unit)
  end subroutine flush_files

  !> Nonlinear convergence condition - simple and experimental look
  !> for the averaged differential of the summed averaged heat flux to
  !> drop below a threshold
  subroutine check_nonlin_convergence(istep, heat_flux, exit)
    use diagnostics_nonlinear_convergence, only: check_nonlin_convergence_calc
    logical, intent(inout) :: exit
    integer, intent(in) :: istep
    real, intent(in) :: heat_flux
    call check_nonlin_convergence_calc(istep, nwrite, heat_flux, exit, exit_when_converged, &
         conv_nstep_av, conv_nsteps_converged, conv_test_multiplier, conv_min_step, &
         conv_max_step)
  end subroutine check_nonlin_convergence

  !> Calculate some heating diagnostics, and update their time history
  subroutine heating(istep, navg, h, hk, h_hist, hk_hist)
    use mp, only: proc0
    use fields_arrays, only: phi, apar, bpar, phinew, aparnew, bparnew
    use species, only: nspec, spec
    use kt_grids, only: naky, ntheta0, aky, akx
    use theta_grid, only: ntgrid, delthet, jacob
    use gs2_heating, only: heating_diagnostics, avg_h, avg_hk, zero_htype, get_heat
    use collisions, only: collisions_heating => heating, c_rate
    implicit none
    integer, intent (in) :: istep, navg
    !> Heating diagnostics summed over space at the current timestep
    type (heating_diagnostics), intent(in out) :: h
    type (heating_diagnostics), dimension (0:), intent(in out) :: h_hist
    !> Heating diagnostics as a function of \((\theta, k_y)\) at the current timestep
    type (heating_diagnostics), dimension(:,:), intent(in out) :: hk
    type (heating_diagnostics), dimension(:,:,0:), intent(in out) :: hk_hist
    real, dimension(-ntgrid:ntgrid) :: wgt
    real :: fac
    integer :: is, ik, it, ig

    !Zero out variables for heating diagnostics
    call zero_htype(h)
    call zero_htype(hk)

    if (proc0 .and. collisions_heating .and. allocated(c_rate)) then
       !GGH NOTE: Here wgt is 1/(2*ntgrid+1)
       wgt = delthet*jacob
       wgt = wgt/sum(wgt)

       do is = 1, nspec
          do ik = 1, naky
             fac = 0.5
             if (aky(ik) < epsilon(0.)) fac = 1.0
             do it = 1, ntheta0
                if (aky(ik) < epsilon(0.0) .and. abs(akx(it)) < epsilon(0.0)) cycle
                do ig = -ntgrid, ntgrid

                   !Sum heating by k over all z points (ig)
                   hk(it, ik) % collisions(is) = hk(it, ik) % collisions(is) &
                        + real(c_rate(ig,it,ik,is,1))*fac*wgt(ig)*spec(is)%temp*spec(is)%dens

                   hk(it, ik) % hypercoll(is) = hk(it, ik) % hypercoll(is) &
                        + real(c_rate(ig,it,ik,is,2))*fac*wgt(ig)*spec(is)%temp*spec(is)%dens

                   hk(it, ik) % imp_colls(is) = hk(it, ik) % imp_colls(is) &
                        + real(c_rate(ig,it,ik,is,3))*fac*wgt(ig)*spec(is)%temp*spec(is)%dens

                end do
                h % collisions(is) = h % collisions(is) + hk(it, ik) % collisions(is)
                h % hypercoll(is)  = h % hypercoll(is)  + hk(it, ik) % hypercoll(is)
                h % imp_colls(is)  = h % imp_colls(is)  + hk(it, ik) % imp_colls(is)
             end do
          end do
       end do
    end if

    ! We may be able to skip this work as well if heating is false.
    call get_heat (h, hk, phi, apar, bpar, phinew, aparnew, bparnew)

    call avg_h(h, h_hist, istep, navg)
    call avg_hk(hk, hk_hist, istep, navg)
  end subroutine heating

  !> Calculate the time-averaged antenna current \(j_{ext} = k_\perp^2 A_{antenna}\)
  subroutine calc_jext (istep, j_ext)
    use mp, only: proc0
    use antenna, only: get_jext
    implicit none
    !> Current simulation timestep
    integer, intent (in) :: istep
    !Shouldn't really need intent in here but it's beacuse we zero it before calling calc_jext
    real, dimension(:,:), intent(in out) ::  j_ext

    integer :: i
    if (.not. (proc0 .and. navg > 1)) return
    call get_jext(j_ext)
    ! This looks a little odd as we ignore the first step
    if (istep > 1) j_ext_hist(:, :, mod(istep, navg)) = j_ext(:, :)

    !Use average of history
    if (istep >= navg) then
       j_ext = 0.
       do i = 0, navg - 1
          j_ext = j_ext + j_ext_hist(:, :, i) / real(navg)
       end do
    end if
  end subroutine calc_jext

  !> A linear estimate of the diffusivity, primarily used testing
  pure real function diffusivity()
    use kt_grids, only: ntheta0, naky, kperp2
    use theta_grid, only: grho
    use warning_helpers, only: is_zero
    real, dimension(ntheta0, naky) :: diffusivity_by_k
    integer  :: ik, it
    diffusivity_by_k = 0.0
    do ik = 1, naky
       do it = 1, ntheta0
          if (is_zero(kperp2(igomega,it,ik))) cycle
          diffusivity_by_k(it,ik) = 2.0 &
               * max(aimag(omegaavg(it,ik)), 0.0) / kperp2(igomega, it, ik)
       end do
    end do
    diffusivity = maxval(diffusivity_by_k) * grho(igomega)
  end function diffusivity

  !> Calculates <|field|**2 kperp2>_theta / <|field|**2>_theta.
  !> Useful for simple quasilinear metric
  pure function get_field_weighted_average_kperp2(field) result(average)
    use kt_grids, only: kperp2, naky, ntheta0
    use theta_grid, only: ntgrid, field_line_average
    implicit none
    complex, dimension(-ntgrid:ntgrid, ntheta0, naky), intent(in) :: field
    real, dimension(ntheta0, naky) :: average
    real, dimension(-ntgrid:ntgrid, ntheta0, naky) :: field_squared
    integer :: it, ik
    field_squared = abs(field) * abs(field)
    average = 0.0
    do ik = 1, naky
       do it = 1, ntheta0
          if (maxval(abs(field_squared(:, it, ik))) <= epsilon(0.0)) cycle
          average(it, ik) = field_line_average(field_squared(:, it, ik) * &
               kperp2(:, it, ik)) / field_line_average(field_squared(:, it, ik))
       end do
    end do
  end function get_field_weighted_average_kperp2

  !> Calculates a simple gamma/kperp2 QL flux metric
  function calculate_simple_quasilinear_flux_metric_by_k(growth_rates) result(ql_metric)
    use run_parameters, only: has_phi, has_apar, has_bpar
    use kt_grids, only: naky, ntheta0, aky
    use fields_arrays, only: phinew, aparnew, bparnew
    use warning_helpers, only: is_zero
    implicit none
    real, dimension(ntheta0, naky), intent(in) :: growth_rates
    real, dimension(ntheta0, naky) :: limited_growth_rates, average
    real, dimension(ntheta0, naky) :: ql_metric, normalising_field
    integer :: ik

    ! Initialise flux to zero
    ql_metric = 0.0

    ! Decide on the normalising field
    if (has_phi) then
       normalising_field = maxval(abs(phinew), dim = 1)
    else if (has_apar) then
       normalising_field = maxval(abs(aparnew), dim = 1)
    else if (has_bpar) then
       normalising_field = maxval(abs(bparnew), dim = 1)
    else
       ! If we get here then no fields are active so the QL flux
       ! is zero and we can exit
       return
    end if

    where (normalising_field > 0)
       normalising_field = 1.0 / normalising_field
    end where

    if (has_phi) then
       average = get_field_weighted_average_kperp2(phinew)
       where(average > 0.0)
          ql_metric = maxval(abs(phinew), dim = 1) * normalising_field / average
       end where
    end if

    if (has_apar) then
       average = get_field_weighted_average_kperp2(aparnew)
       where(average > 0.0)
          ql_metric = ql_metric + &
               maxval(abs(aparnew), dim = 1) * normalising_field / average
       end where
    end if

    if (has_bpar) then
       average = get_field_weighted_average_kperp2(bparnew)
       where(average > 0.0)
          ql_metric = ql_metric + &
               maxval(abs(bparnew), dim = 1) * normalising_field / average
       end where
    end if

    ! Limit the growth rate to positive values
    where (growth_rates > 0)
       limited_growth_rates = growth_rates
    elsewhere
       limited_growth_rates = 0.0
    end where

    ! MG: Avoid spurious contribution from the zonal mode
    do ik = 1, naky
       if (is_zero(aky(ik))) ql_metric(:, ik) = 0.0
    end do

    ql_metric = ql_metric * limited_growth_rates
  end function calculate_simple_quasilinear_flux_metric_by_k

  !> Calculates the instantaneous omega from chinew = chi * exp(-i * omega * dt)
  !> with chi = phi + apar + bpar. This gives omega = log(chinew/chi) * i / dt.
  pure function  calculate_instantaneous_omega(ig, tolerance) result(omega_inst)
    use kt_grids, only: naky, ntheta0
    use fields_arrays, only: phi, apar, bpar, phinew, aparnew, bparnew
    use gs2_time, only: code_dt
    use constants, only: zi
    use optionals, only: get_option_with_default
    use run_parameters, only: has_phi, has_apar, has_bpar
    use warning_helpers, only: is_zero
    implicit none
    integer, intent(in), optional :: ig
    real, intent(in), optional :: tolerance
    complex, dimension(ntheta0, naky) :: omega_inst, field_sum, field_sum_new
    integer :: ig_use
    real :: tol_use, too_small

    ! Determine which theta grid point to evaluate the fields at
    ig_use = get_option_with_default(ig, igomega)

    ! Construct the field sums for current and old fields.
    ! We always allocate the fields (for now) so we could just
    ! unconditionally sum. This could make the code clearer
    ! without really adding _much_ overhead.
    if (has_phi) then
       field_sum = phi(ig_use, :, :)
       field_sum_new = phinew(ig_use, :, :)
    else
       field_sum = 0.0
       field_sum_new = 0.0
    end if

    if (has_apar) then
       field_sum = field_sum + apar(ig_use, :, :)
       field_sum_new = field_sum_new + aparnew(ig_use, :, :)
    end if

    if (has_bpar) then
       field_sum = field_sum + bpar(ig_use, :, :)
       field_sum_new = field_sum_new + bparnew(ig_use, :, :)
    end if

    ! Determine what tolerance to use
    tol_use = get_option_with_default(tolerance, omegatol)

    ! Decide what size float we consider too small to divide by
    ! Note this uses tiny rather than epsilon in order to allow
    ! us to track modes with low amplitudes.
    if (is_zero(tol_use)) then
       too_small = tiny(0.0)
    else
       too_small = tiny(0.0) * 1000. / abs(tol_use)
    end if

    ! Use where to guard against dividing by tool small a number
    ! Note we don't divide by field_sum_new so we probably don't need to
    ! check that here, just field_sum.
    where (abs(field_sum) < too_small .or. abs(field_sum_new) < too_small)
       omega_inst = 0.0
    elsewhere
       omega_inst = log(field_sum_new / field_sum) * zi / code_dt
    end where

  end function calculate_instantaneous_omega

  !> Calculate the time-averaged complex frequency, check convergence
  !> criterion, and numerical instability criterion.
  !>
  !> Time average is over previous [[gs2_diagnostics_knobs:navg]] timesteps
  subroutine get_omegaavg (istep, exit, omegaavg, debopt)
    use kt_grids, only: ntheta0, naky
    use gs2_time, only: code_dt, user_time
    use optionals, only: get_option_with_default
    implicit none
    !> The current timestep
    integer, intent (in) :: istep
    !> Should the simulation exit?
    !> FIXME: This could be `intent(out)` only
    logical, intent (in out) :: exit
    !> The time-averaged complex frequency
    complex, dimension (:,:), intent (out) :: omegaavg
    !> If true, write some debug messages
    logical, intent (in), optional :: debopt
    logical :: debug
    debug = get_option_with_default(debopt, .false.)
    if (debug) write(6,*) "get_omeaavg: allocate domega"
    if (.not. allocated(domega)) allocate(domega(navg,ntheta0,naky))
    if (debug) write(6,*) "get_omeaavg: start"
    if (debug) write(6,*) "get_omeaavg: set omegahist"

    ! Get and store the current instantaneous omega
    omegahist(mod(istep, navg), :, :) = calculate_instantaneous_omega(igomega, omegatol)

    if (debug) write(6,*) "get_omeaavg: set omegahist at istep = 0"
    !During initialisation fieldnew==field but floating error can lead to finite omegahist
    !Force omegahist=0 here to avoid erroneous values.
    !Could think about forcing omegahist=0 where abs(omegahist)<tol
    if(istep==0) omegahist(:,:,:)=0.0

    if (debug) write(6,*) "get_omeaavg: set omegaavg"
    omegaavg = sum(omegahist/real(navg),dim=1)
    if (debug) write(6,*) "get_omegaavg: omegaavg=",omegaavg
    if ((istep > navg) .and.(user_time >= omega_checks_start)) then
       domega = spread(omegaavg,1,navg) - omegahist
       if (all(sqrt(sum(abs(domega)**2/real(navg),dim=1)) &
            .le. min(abs(omegaavg),1.0)*omegatol)) &
            then
          if (write_ascii) write (out_unit, "('*** omega converged')")
          if (exit_when_converged) exit = .true.
       end if

       if (any(abs(omegaavg)*code_dt > omegatinst)) then
          if (write_ascii) write (out_unit, "('*** numerical instability detected')") 
          exit = .true.
       end if

       if (all(aimag(omegaavg) < damped_threshold)) then
          if (write_ascii) write (out_unit, "('*** all modes strongly damped')")
          exit = .true.
       end if
    end if
    if (debug) write(6,*) "get_omegaavg: done"
  end subroutine get_omegaavg

  !> FIXME : Add documentation
  !> @note: Should look at moving this and weighted_ky_sum to volume_averages
  subroutine get_volume_average (f, favg)
    implicit none
    real, dimension (:,:), intent (in) :: f
    real, intent (out) :: favg
    favg = sum(weighted_ky_sum(f))
  end subroutine get_volume_average

  !> Sum a {kx,ky} array over ky accounting for non-uniform weighting arising
  !> from gs2's non-standard Fourier coefficients:
  !> ky=0 modes have correct amplitudes; rest must be scaled
  !> note contrast with scaling factors in FFT routines.
  !>
  !>  fac values here arise because gs2's Fourier coefficients, F_k^gs2, not standard form:
  !>          i.e. f(x) = f_k e^(i k.x)
  !>  With standard Fourier coeffs in gs2, we would instead need:  fac=2.0 for ky > 0
  !>      (fac=2.0 would account ky<0 contributions, not stored due to reality condition)
  pure function weighted_ky_sum(f) result(ky_sum)
    use kt_grids, only: naky, ntheta0, aky
    use warning_helpers, only: is_zero
    implicit none
    real, dimension(:, :), intent (in) :: f
    real, dimension(ntheta0) :: ky_sum
    real :: fac
    integer :: ik
    ky_sum = 0.
    do ik = 1, naky
       fac = 0.5
       if (is_zero(aky(ik))) fac = 1.0
       ky_sum = ky_sum + f(:, ik) * fac
    end do
  end function weighted_ky_sum

  !> Calculate the correlation function over the extended domain
  subroutine correlation_extend (cfnc, phi2extend, ntg_extend)
    use fields_arrays, only: phinew
    use theta_grid, only: ntgrid, jacob, delthet, ntgrid
    use kt_grids, only: ntheta0, naky, jtwist
    implicit none
    real, dimension (:,:,:), intent (out) :: phi2extend
    complex, dimension (:,:,:), intent (out) :: cfnc
    integer, intent(in) :: ntg_extend
    integer :: ig, it, ik, im, igmod, itshift, nconnect, offset
    real, dimension (:), allocatable :: dl_over_b
    complex, dimension (:,:,:), allocatable :: phiextend
    complex, dimension (:,:), allocatable :: phir

    allocate (dl_over_b(2*ntgrid+1))
    dl_over_b = delthet * jacob

    allocate (phiextend(ntg_extend,ntheta0,naky)) ; phiextend = 0.0
    allocate (phir(-ntgrid:ntgrid,ntheta0))

    cfnc = 0.0 ; phiextend = 0.0
    offset = ((ntheta0-1)/jtwist)*(2*ntgrid+1)/2
    call reorder_kx (phinew(:, :, 1), phir(:, :))
    phiextend(offset+1:offset+(2*ntgrid+1),:,1) = phir
    do it = 1, ntheta0
       do im = 1, 2*ntgrid+1
          do ig = im+offset, offset+(2*ntgrid+1)
             igmod = mod(ig-offset-1,2*ntgrid+1)+1
             cfnc(im,it,1) = cfnc(im,it,1) &
                  + phiextend(ig,it,1)*conjg(phiextend(ig-im+1,it,1)) &
                  * dl_over_b(igmod)
          end do
       end do
       if (abs(cfnc(1, it, 1)) > epsilon(0.0)) then
         cfnc(:, it, 1) = cfnc(:, it, 1) / cfnc(1, it, 1)
       end if
    end do

    do ik = 2, naky
       call reorder_kx (phinew(:, :, ik), phir(:, :))
       ! shift in kx due to parallel boundary condition
       ! also the number of independent theta0s
       itshift = jtwist*(ik-1)
       do it = 1, min(itshift,ntheta0)
          ! number of connections between kx's
          nconnect = (ntheta0-it)/itshift
          ! shift of theta index to account for fact that not all ky's
          ! have same length in extended theta
          offset = (2*ntgrid+1)*((ntheta0-1)/jtwist-nconnect)/2
          do ig = 1, nconnect+1
             phiextend(offset+(ig-1)*(2*ntgrid+1)+1:offset+ig*(2*ntgrid+1),it,ik) &
                  = phir(:,ntheta0-it-(ig-1)*itshift+1)
          end do
          do im = 1, (2*ntgrid+1)*(nconnect+1)
             do ig = im+offset, offset+(2*ntgrid+1)*(nconnect+1)
                igmod = mod(ig-offset-1,2*ntgrid+1)+1
                cfnc(im,it,ik) = cfnc(im,it,ik) &
                     + phiextend(ig,it,ik)*conjg(phiextend(ig-im+1,it,ik)) &
                     * dl_over_b(igmod)
             end do
          end do
          cfnc(:,it,ik) = cfnc(:,it,ik) / cfnc(1,it,ik)
       end do
    end do

    phi2extend = real(phiextend*conjg(phiextend))

    deallocate (dl_over_b, phir, phiextend)
  end subroutine correlation_extend

  !> Go from kx_ind = [0, 1, ..., N, -N, ...., -1]
  !> to [-N, ....,-1,0,1,....,N]
  subroutine reorder_kx (unord, ord)
    use kt_grids, only: ntheta0
    implicit none
    complex, dimension (:, :), intent (in) :: unord
    complex, dimension (:, :), intent (out) :: ord
    ord = cshift(unord, -ntheta0 / 2, dim = 2)
  end subroutine reorder_kx

  !> Calculate the correlation function on the physical domain
  subroutine correlation (cfnc_2pi)
    use kt_grids, only: naky, ntheta0
    use theta_grid, only: ntgrid
    use fields_arrays, only: phinew
    implicit none
    complex, dimension (-ntgrid:,:), intent (out) :: cfnc_2pi
    integer :: ik, it, ig
    real :: fac

    cfnc_2pi = 0.0

    do ik = 1, naky
       if (ik==1) then
          fac = 1.0
       else
          fac = 0.5
       end if
       do it = 1, ntheta0
          do ig = -ntgrid, ntgrid
             cfnc_2pi(ig,ik) = cfnc_2pi(ig,ik) + phinew(0,it,ik)*conjg(phinew(ig,it,ik))*fac
          end do
       end do
    end do

  end subroutine correlation

  !> Write \(g(v_\parallel,v_\perp)\) at `ik=it=is=1, ig=0` to text file `<runname>.dist`
  subroutine do_write_f (unit)
    use mp, only: proc0, send, receive
    use gs2_layouts, only: g_lo, ik_idx, it_idx, is_idx, il_idx, ie_idx, idx_local, proc_id
    use le_grids, only: al, energy, forbid, negrid, nlambda, speed
    use theta_grid, only: bmag
    use dist_fn_arrays, only: gnew
    use species, only: nspec
    integer, intent(in) :: unit
    integer :: iglo, ik, it, is
    integer :: ie, il, ig
    real :: vpa, vpe
    complex, dimension(2) :: gtmp
    real, dimension (:,:), allocatable, save :: xpts
    real, dimension (:), allocatable, save :: ypts
    !> Change argument of bmag depending on which theta you want to write out
    integer, parameter :: ig_index = 0
    if (.not. allocated(xpts)) then
       allocate(xpts(negrid,nspec))
       allocate(ypts(nlambda))

       ! should really just get rid of xpts and ypts
       xpts = speed
       ypts = 0.0

       do il=1,nlambda
          ypts(il) = sqrt(max(0.0,1.0-al(il)*bmag(ig_index)))
       end do

       if (proc0) then
          write(unit, *) negrid*nlambda
       end if
    endif

    do iglo = g_lo%llim_world, g_lo%ulim_world
       ik = ik_idx(g_lo, iglo) ; if (ik /= 1) cycle
       it = it_idx(g_lo, iglo) ; if (it /= 1) cycle
       is = is_idx(g_lo, iglo) ; if (is /= 1) cycle
       ie = ie_idx(g_lo, iglo)
       ig = ig_index
       il = il_idx(g_lo, iglo)
       if (idx_local (g_lo, ik, it, il, ie, is)) then
          if (proc0) then
             gtmp = gnew(ig, :, iglo)
          else
             call send (gnew(ig,:,iglo), 0)
          endif
       else if (proc0) then
          call receive (gtmp, proc_id(g_lo, iglo))
       endif
       if (proc0) then
          if (.not. forbid(ig, il)) then
             vpa = sqrt(energy(ie,is)*max(0.0, 1.0-al(il)*bmag(ig)))
             vpe = sqrt(energy(ie,is)*al(il)*bmag(ig))
             write (unit, "(8(1x,e13.6))") vpa, vpe, energy(ie,is), al(il), &
                  xpts(ie,is), ypts(il), real(gtmp(1)), real(gtmp(2))
          end if
       end if
    end do
    if (proc0) write (unit, *)
  end subroutine do_write_f

  !> Write distribution function (\(g\) or possibly \(f\)?) in real space to text file
  !> "<runname>.yxdist"
  subroutine do_write_fyx (unit, phi, bpar)
    use mp, only: proc0, send, receive, barrier
    use gs2_layouts, only: il_idx, ig_idx, ik_idx, it_idx, is_idx, isign_idx, ie_idx
    use gs2_layouts, only: idx_local, proc_id, yxf_lo, accelx_lo, g_lo
    use le_grids, only: al, energy=>energy_maxw, forbid, negrid, nlambda
    use theta_grid, only: bmag, ntgrid
    use dist_fn_arrays, only: gnew, g_adjust, g_work, to_g_gs2, from_g_gs2
    use nonlinear_terms, only: accelerated
    use gs2_transforms, only: transform2
    use array_utils, only: zero_array, copy
#ifdef SHMEM
    use shm_mpi3, only: shm_alloc, shm_free
#endif
    integer, intent(in) :: unit
    !> Electrostatic potential and parallel magnetic field
    complex, dimension (-ntgrid:,:,:), intent (in) :: phi, bpar
    real, dimension (:,:), allocatable :: grs, gzf
#ifndef SHMEM
    real, dimension (:,:,:), allocatable :: agrs, agzf
#else
    real, dimension (:,:,:), pointer, contiguous :: agrs => null(), agzf => null()
    complex, dimension(:,:,:), pointer, contiguous :: g0_ptr => null()
#endif
    real, dimension (:), allocatable :: agp0, agp0zf
    real :: gp0, gp0zf
    integer :: ig, it, ik, il, ie, is, iyxlo, isign, iaclo, iglo, acc
    logical, save :: first = .true.

    if (accelerated) then
#ifndef SHMEM
       allocate (agrs(2*ntgrid+1, 2, accelx_lo%llim_proc:accelx_lo%ulim_alloc))
       allocate (agzf(2*ntgrid+1, 2, accelx_lo%llim_proc:accelx_lo%ulim_alloc))
#else
       call shm_alloc(agrs, [1, 2*ntgrid+1, 1, 2, accelx_lo%llim_proc, accelx_lo%ulim_alloc])
       call shm_alloc(agzf, [1, 2*ntgrid+1, 1, 2, accelx_lo%llim_proc, accelx_lo%ulim_alloc])
       call shm_alloc(g0_ptr, [1, 2*ntgrid+1, 1, 2, accelx_lo%llim_proc, accelx_lo%ulim_alloc])
#endif
       allocate (agp0(2), agp0zf(2))
       call zero_array(agrs); call zero_array(agzf); agp0 = 0.0; agp0zf = 0.0
    else
#ifdef SHMEM
       g0_ptr => g_work
#endif
       allocate (grs(yxf_lo%ny, yxf_lo%llim_proc:yxf_lo%ulim_alloc))
       allocate (gzf(yxf_lo%ny, yxf_lo%llim_proc:yxf_lo%ulim_alloc))
       call zero_array(grs); call zero_array(gzf); gp0 = 0.0; gp0zf = 0.0
    end if

    if (first) then
       if (proc0) then
          acc = 0
          if (accelerated) acc = 1
          write(unit,*) 2*negrid*nlambda, bmag(0), acc
       end if

       first = .false.
    end if

    call g_adjust (gnew, phi, bpar, direction = from_g_gs2)

#ifndef SHMEM
    call copy(gnew, g_work)
#else
    g0_ptr = gnew
#endif

#ifndef SHMEM
    if (accelerated) then
       call transform2 (g_work, agrs)
    else
       call transform2 (g_work, grs)
    end if

    call zero_array(g_work)
    do iglo=g_lo%llim_proc, g_lo%ulim_proc
       ik = ik_idx(g_lo,iglo)
       if (ik == 1) g_work(:,:,iglo) = gnew(:,:,iglo)
    end do

    call g_adjust (gnew, phi, bpar, direction = to_g_gs2)

    if (accelerated) then
       call transform2 (g_work, agzf)
    else
       call transform2 (g_work, gzf)
    end if
#else
    if (accelerated) then
       call transform2 (g0_ptr, agrs)
    else
       call transform2 (g0_ptr, grs)
    end if

    g0_ptr = 0.0
    do iglo=g_lo%llim_proc, g_lo%ulim_proc
       ik = ik_idx(g_lo,iglo)
       if (ik == 1) g0_ptr(:,:,iglo) = gnew(:,:,iglo)
    end do

    call g_adjust (gnew, phi, bpar, direction = to_g_gs2)

    if (accelerated) then
       call transform2 (g0_ptr, agzf)
    else
       call transform2 (g0_ptr, gzf)
    end if
#endif

    if (accelerated) then
       do iaclo=accelx_lo%llim_world, accelx_lo%ulim_world
          it = it_idx(accelx_lo, iaclo)
          ik = ik_idx(accelx_lo, iaclo)
          if (it == 1 .and. ik == 1) then
             il = il_idx(accelx_lo, iaclo)
             ig = 0
             if (.not. forbid(ig,il)) then
                ie = ie_idx(accelx_lo, iaclo)
                is = is_idx(accelx_lo, iaclo)

                if (proc0) then
                   if (.not. idx_local(accelx_lo, ik, it, il, ie, is)) then
                      call receive (agp0, proc_id(accelx_lo, iaclo))
                      call receive (agp0zf, proc_id(accelx_lo, iaclo))
                   else
                      agp0 = agrs(ig+ntgrid+1, :, iaclo)
                      agp0zf = agzf(ig+ntgrid+1, :, iaclo)
                   end if
                else if (idx_local(accelx_lo, ik, it, il, ie, is)) then
                   call send (agrs(ig+ntgrid+1, :, iaclo), 0)
                   call send (agzf(ig+ntgrid+1, :, iaclo), 0)
                end if

                if (proc0) then
                   write (unit, "(6(1x,e13.6))") energy(ie), al(il), &
                        agp0(1), agp0(2), agp0zf(1), agp0zf(2)
                end if
             end if
          end if
          call barrier
       end do
       deallocate(agp0, agp0zf)
#ifndef SHMEM
       deallocate(agrs, agzf)
#else
       call shm_free(agrs)
       call shm_free(agzf)
       call shm_free(g0_ptr)
#endif
    else
       do iyxlo=yxf_lo%llim_world, yxf_lo%ulim_world

          ig = ig_idx(yxf_lo, iyxlo)
          it = it_idx(yxf_lo, iyxlo)
          if (ig == 0 .and. it == 1) then
             il = il_idx(yxf_lo, iyxlo)
             if (.not. forbid(ig,il)) then
                ik = 1
                ie = ie_idx(yxf_lo, iyxlo)
                is = is_idx(yxf_lo, iyxlo)
                isign = isign_idx(yxf_lo, iyxlo)

                if (proc0) then
                   if (.not. idx_local(yxf_lo, ig, isign, it, il, ie, is)) then
                      call receive (gp0, proc_id(yxf_lo, iyxlo))
                      call receive (gp0zf, proc_id(yxf_lo, iyxlo))
                   else
                      gp0 = grs(ik, iyxlo)
                      gp0zf = gzf(ik, iyxlo)
                   end if
                else if (idx_local(yxf_lo, ig, isign, it, il, ie, is)) then
                   call send (grs(ik, iyxlo), 0)
                   call send (gzf(ik, iyxlo), 0)
                end if

                if (proc0) then
                   write (unit, "(4(1x,e13.6),i8)") energy(ie), al(il), &
                        gp0, gp0zf, isign
                end if
             end if
          end if
          call barrier
       end do
       deallocate (grs, gzf)
    end if

    if (proc0) flush (unit)

    if (proc0) then
       write(unit,*)
    end if

  end subroutine do_write_fyx
end module gs2_diagnostics

gyrokinetics / gs2 / 2078172070

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous