gs2_diagnostics.fpp Source File


Contents

Source Code


Source Code

!> 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_base_config, only: diagnostics_base_config_type
  use gs2_io, only: get_netcdf_file_id

  ! what about boundary condition contributions?  In the presence 
  ! of magnetic shear, there are not always zero amplitudes at the ends
  ! of the supercells.

  implicit none

  private

  public :: read_parameters
  public :: init_gs2_diagnostics
  public :: finish_gs2_diagnostics
  public :: check_gs2_diagnostics
  public :: wnml_gs2_diagnostics
  public :: loop_diagnostics
  public :: ensemble_average
  public :: reset_init
  public :: pflux_avg, qflux_avg, heat_avg, vflux_avg, start_time
  public :: diffusivity, calculate_simple_quasilinear_flux_metric_by_k
  public :: get_omegaavg, get_volume_average, calculate_instantaneous_omega

  public :: gs2_diagnostics_config_type
  public :: set_gs2_diagnostics_config
  public :: get_gs2_diagnostics_config
  public :: apply_old_defaults

  ! Export some functions to be used in the new diagnostics
  public :: save_restart_dist_fn
  public :: check_restart_file_writeable
  public :: do_dump_fields_periodically, do_write_parity
  public :: do_write_heating, do_write_nl_flux_dist
  public :: do_write_correlation_extend, do_write_correlation
  public :: do_write_geom, do_write_symmetry, do_write_pflux_sym

  interface get_vol_average
     module procedure get_vol_average_one, get_vol_average_all
  end interface get_vol_average

  interface get_vol_int
     module procedure get_vol_int_one, get_vol_int_all
  end interface get_vol_int

  !> Calculate the field-line average
  interface get_fldline_avg
     module procedure get_fldline_avg_r, get_fldline_avg_c
  end interface get_fldline_avg

  real, public :: omegatol, omegatinst
  logical, public :: print_line, print_flux_line
  logical, public :: print_summary, write_line, write_flux_line
  logical, public :: write_omega, write_omavg, write_ascii
  logical, public :: write_gs, write_ql_metric
  logical, public :: write_g, write_gyx
  logical, public :: write_eigenfunc, write_fields, write_final_fields, write_final_antot
  logical, public :: write_final_moments, write_avg_moments, write_parity
  logical, public :: write_moments, ob_midplane, write_final_db
  logical, public :: write_full_moments_notgc, write_cross_phase = .false.
  logical, public :: write_final_epar, write_kpar
  logical, public :: write_heating, write_lorentzian
  logical, public :: write_fluxes, write_fluxes_by_mode
  logical, public :: write_verr, write_cerr, write_max_verr
  logical, public :: exit_when_converged
  logical, public :: use_nonlin_convergence
  logical, public :: dump_check1, dump_check2
  logical, public :: dump_fields_periodically, make_movie
  logical, public :: save_for_restart
  logical, public :: save_distfn
  logical, public :: save_glo_info_and_grids
  logical, public :: save_velocities
  logical, public :: write_symmetry, write_correlation_extend, write_correlation
  logical, public :: write_pflux_sym, write_nl_flux_dist, write_pflux_tormom
  logical :: file_safety_check
  integer, public :: nwrite, igomega, nmovie
  integer, public :: navg, nsave, nwrite_mult, nc_sync_freq

  logical, public :: write_phi_over_time, write_apar_over_time, write_bpar_over_time !EGH

  logical :: append_old
  logical :: write_jext=.true.

  !> Variables for convergence condition testing
  integer :: trin_istep = 0
  integer :: conv_isteps_converged = 0
  real, allocatable, dimension(:) :: conv_heat
  real :: heat_sum_av = 0, heat_av = 0, heat_av_test = 0


  ! 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, dump_any
  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
  integer :: phase_unit
  integer :: dump_check1_unit, dump_check2_unit
  integer :: res_unit, res_unit2, parity_unit
  integer :: conv_nstep_av
  integer :: conv_min_step
  integer :: 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, dimension (:,:,:), allocatable ::  pflux_tormom

  ! (ntheta0,naky,nspec)

  real :: start_time = 0.0
  real, dimension (:), allocatable :: pflux_avg, qflux_avg, heat_avg, vflux_avg
  real :: conv_test_multiplier

  integer :: ntg_extend, nth0_extend
  integer :: nout = 1
  integer :: nout_movie = 1
  integer :: nout_big = 1
  complex :: wtmp_old = 0.
  !> Does the `gs2_diagnostics_knobs` namelist exist in the input file?
  logical :: exist

  type, extends(diagnostics_base_config_type) :: gs2_diagnostics_config_type
  end type gs2_diagnostics_config_type

  type(gs2_diagnostics_config_type) :: gs2_diagnostics_config
contains

  !> Write the diagnostics namelist to file
  subroutine wnml_gs2_diagnostics(unit)
    implicit none
    !> Unit of an open file to write to
    integer, intent(in) :: unit
    if (.not.exist) return
    write (unit, *)
    write (unit, fmt="(' &',a)") "gs2_diagnostics_knobs"
    write (unit, fmt="(' save_for_restart = ',L1)") save_for_restart
    write (unit, fmt="(' save_distfn = ',L1)") save_distfn
    write (unit, fmt="(' save_many = ',L1)") save_many
    write (unit, fmt="(' file_safety_check = ',L1)") file_safety_check
    write (unit, fmt="(' print_line = ',L1)") print_line 
    write (unit, fmt="(' write_line = ',L1)") write_line
    write (unit, fmt="(' print_flux_line = ',L1)") print_flux_line
    write (unit, fmt="(' write_flux_line = ',L1)") write_flux_line
    write (unit, fmt="(' nmovie = ',i6)") nmovie
    write (unit, fmt="(' nwrite_mult = ',i6)") nwrite_mult
    write (unit, fmt="(' nwrite = ',i6)") nwrite
    write (unit, fmt="(' nsave = ',i6)") nsave
    write (unit, fmt="(' navg = ',i6)") navg
    write (unit, fmt="(' omegatol = ',e17.10)") omegatol
    write (unit, fmt="(' omegatinst = ',e17.10)") omegatinst
    ! should be legal -- not checked yet
    if (igomega /= 0) write (unit, fmt="(' igomega = ',i6)") igomega  

    if (write_ascii) then
       write (unit, fmt="(' write_ascii = ',L1)") write_ascii
       write (unit, fmt="(' write_omega = ',L1)") write_omega
       write (unit, fmt="(' write_omavg = ',L1)") write_omavg
    end if
    write (unit, fmt="(' write_heating = ',L1)") write_heating
    write (unit, fmt="(' write_lorentzian = ',L1)") write_lorentzian
    write (unit, fmt="(' write_eigenfunc = ',L1)") write_eigenfunc
    write (unit, fmt="(' write_final_fields = ',L1)") write_final_fields
    write (unit, fmt="(' write_final_epar = ',L1)") write_final_epar
    write (unit, fmt="(' write_final_db = ',L1)") write_final_db
    write (unit, fmt="(' write_final_moments = ',L1)") write_final_moments
    write (unit, fmt="(' write_final_antot = ',L1)") write_final_antot
    write (unit, fmt="(' write_fluxes = ',L1)") write_fluxes
    write (unit, fmt="(' write_fluxes_by_mode = ',L1)") write_fluxes_by_mode
    write (unit, fmt="(' exit_when_converged = ',L1)") exit_when_converged
    write (unit, fmt="(' use_nonlin_convergence = ',L1)") use_nonlin_convergence
    if (write_avg_moments) write (unit, fmt="(' write_avg_moments = ',L1)") write_avg_moments
    if (dump_check1) write (unit, fmt="(' dump_check1 = ',L1)") dump_check1
    if (dump_check2) write (unit, fmt="(' dump_check2 = ',L1)") dump_check2
    if (dump_fields_periodically) &
         write (unit, fmt="(' dump_fields_periodically = ',L1)") dump_fields_periodically
    if (make_movie) &
         write (unit, fmt="(' make_movie = ',L1)") make_movie

    write (unit, fmt="(' /')")       
  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: nonlinear_mode_switch, nonlinear_mode_on
    use dist_fn, only : def_parity, even 
    use kt_grids, only : gridopt_switch, gridopt_box
    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_avg_moments) then
       if (gridopt_switch /= gridopt_box) then
          write (report_unit, *) 
          write (report_unit, fmt="('################# WARNING #######################')")
          write (report_unit, fmt="('write_avg_moments = T:          Ignored unless grid_option=box')")
          write (report_unit, fmt="('################# WARNING #######################')")
          write (report_unit, *) 
       else
          if (write_ascii) then
             write (report_unit, fmt="('write_avg_moments = T:     Flux surface averaged low-order moments of g written to ',a)") &
                  & trim(run_name)//'.moments'
          end if
          write (report_unit, fmt="('write_avg_moments = T:     Flux surface averaged low-order moments of g written to ',a)") &
               & trim(run_name)//'.out.nc'
       end if
    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 (dump_check1) then
       write (report_unit, fmt="('dump_check1 = T:          Field-line avg of Phi written to ',a)") 'dump.check1'
       write (report_unit, fmt="('This option is usually used for Rosenbluth-Hinton calculations.')") 
       write (report_unit, fmt="('THIS IS PROBABLY AN ERROR.')") 
    end if

    if (dump_check2) then
       write (report_unit, fmt="('dump_check2 = T:           Apar(kx, ky, igomega) written to ',a)") trim(run_name)//'.dc2'
       write (report_unit, fmt="('THIS IS PROBABLY AN ERROR.')") 
    end if

    if (dump_fields_periodically) then
       write (report_unit, fmt="('dump_fields_periodically = T:          Phi, Apar, Bpar written to ',a)") 'dump.fields.t=(time)'
       write (report_unit, fmt="('THIS IS PROBABLY AN ERROR.  IT IS EXPENSIVE.')") 
    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 (nonlinear_mode_switch == nonlinear_mode_on) 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 (list, nstep, gs2_diagnostics_config_in, header)
    use theta_grid, only: init_theta_grid
    use kt_grids, only: init_kt_grids, ntheta0, naky
    use run_parameters, only: init_run_parameters
    use species, only: init_species, nspec
    use dist_fn, only: init_dist_fn
    use init_g, only: init_init_g
    use gs2_io, only: init_gs2_io
    use gs2_heating, only: init_htype
    use collisions, only: collision_model_switch, init_lorentz_error
    use mp, only: broadcast, proc0, mp_abort
    use le_grids, only: init_weights
    use gs2_save, only: restart_writable
    use normalisations, only: init_normalisations
    use diagnostics_final_routines, only: init_par_filter
    use standard_header, only: standard_header_type
    implicit none
    !> If true, this is a "list-mode run" and so turn off
    !> [[gs2_diagnostics_knobs:print_flux_line]] and
    !> [[gs2_diagnostics_knobs:print_line]] if set
    logical, intent (in) :: list
    !> Maximum number of steps to take
    integer, intent (in) :: nstep
    !> Input parameters for this module
    type(gs2_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
    integer :: nmovie_tot

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

    if (initialized) return
    initialized = .true.

    call init_normalisations
    call init_theta_grid
    call init_kt_grids
    call init_run_parameters
    call init_species
    call init_init_g
    call init_dist_fn

    call real_init (list, gs2_diagnostics_config_in, header=local_header)

    nmovie_tot = nstep/nmovie

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

    ! 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, &
         write_correlation_extend, &
         write_phi_over_time, write_apar_over_time, write_bpar_over_time, &
         nout, ob_midplane=ob_midplane, header=local_header)

    call broadcast (nout)

    if (write_cerr) then
       if (collision_model_switch == 1 .or. collision_model_switch == 5) then
          call init_lorentz_error
       else
          write_cerr = .false.
       end if
    end if

    if(.not. allocated(pflux_avg)) then
       allocate (pflux_avg(nspec), qflux_avg(nspec), heat_avg(nspec), vflux_avg(nspec))
       pflux_avg = 0.0 ; qflux_avg = 0.0 ; heat_avg = 0.0 ; vflux_avg = 0.0
    endif

    call check_restart_file_writeable(file_safety_check, save_for_restart, save_distfn)

    !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)
    use init_g, only: get_restart_dir
    use gs2_save, only: restart_writable
    use mp, only: proc0, mp_abort, broadcast
    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

    ! 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

    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(is_list_run, gs2_diagnostics_config_in, header)
    use run_parameters, only: fapar
    use file_utils, only: open_output_file, get_unused_unit
    use kt_grids, only: naky, ntheta0
    use mp, only: proc0
    use standard_header, only: standard_header_type
    use diagnostics_fluxes, only: init_diagnostics_fluxes
    implicit none
    !> If true, this is a "list-mode run" and so turn off
    !> [[gs2_diagnostics_knobs:print_flux_line]] and
    !> [[gs2_diagnostics_knobs:print_line]] if set
    logical, intent (in) :: is_list_run
    !> Input parameters for this module
    type(gs2_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(is_list_run, 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")
          if (write_max_verr) call open_output_file (res_unit2, ".vres2")
       end if

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

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

       if (dump_check1 .and. proc0) then
          call get_unused_unit (dump_check1_unit)
          ! We should switch to open_output_file here.
          open (unit=dump_check1_unit, file="dump.check1", status="unknown")
       end if

       if (dump_check2 .and. proc0) then
          call open_output_file (dump_check2_unit, ".dc2")
       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(.not. allocated(conv_heat)) allocate(conv_heat(0:conv_nstep_av/nwrite-1))

    end if

    call init_diagnostics_fluxes()

    if (.not. allocated(omega)) then
       allocate(omega(ntheta0, naky))
       allocate(omegaavg(ntheta0, naky))
       omega=0
       omegaavg=0
    end if

  end subroutine real_init

  !> Set values according to the old diagnostics defaults.
  !>
  !> Will overwrite existing values, unless the config object has
  !> already been initialised, in which case returns without changing
  !> any values
  subroutine apply_old_defaults(gs2_diagnostics_config_in)
    type(gs2_diagnostics_config_type), intent(inout) :: gs2_diagnostics_config_in

    if (gs2_diagnostics_config_in%is_initialised()) return

    gs2_diagnostics_config_in%navg = 100
    gs2_diagnostics_config_in%nmovie = 1000
    gs2_diagnostics_config_in%nwrite = 100
    gs2_diagnostics_config_in%ob_midplane = .true.
    gs2_diagnostics_config_in%omegatinst = 1.0
    gs2_diagnostics_config_in%print_line = .true.
    gs2_diagnostics_config_in%write_correlation = .false.
    gs2_diagnostics_config_in%write_moments = .false.
  end subroutine apply_old_defaults

  !> Read the input parameters for the diagnostics module
  subroutine read_parameters (is_list_run, gs2_diagnostics_config_in)
    use diagnostics_base_config, only: warn_about_nonfunctional_selection
    use file_utils, only: input_unit, input_unit_exist
    use run_parameters, only: has_phi
    use nonlinear_terms, only: nonlin
    use antenna, only: no_driver
    use collisions, only: heating
    use mp, only: proc0
    implicit none
    !> If true, this is a "list-mode run" and so turn off
    !> [[gs2_diagnostics_knobs:print_flux_line]] and
    !> [[gs2_diagnostics_knobs:print_line]] if set
    logical, intent (in) :: is_list_run
    !> Configuration for this module, can be used to set new default values or
    !> avoid reading the input file
    type(gs2_diagnostics_config_type), intent(in), optional :: gs2_diagnostics_config_in
    
    if (present(gs2_diagnostics_config_in)) gs2_diagnostics_config = gs2_diagnostics_config_in

    ! Smart defaults
    if (.not. gs2_diagnostics_config%is_initialised()) then
      if (nonlin) then
        gs2_diagnostics_config%write_fluxes = .true.
        gs2_diagnostics_config%write_fluxes_by_mode = .true.
      end if
      call apply_old_defaults(gs2_diagnostics_config)
    end if

    call gs2_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
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%enable_parallel, "enable_parallel")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%ncheck /= 10, "ncheck")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%serial_netcdf4, "serial_netcdf4")
    call warn_about_nonfunctional_selection(.not. gs2_diagnostics_config%write_any, "write_any")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%write_collisional, "write_collisional")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%write_density_over_time, "write_density_over_time")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%write_ntot_over_time, "write_ntot_over_time")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%write_tperp_over_time, "write_tperp_over_time")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%write_upar_over_time, "write_upar_over_time")
    call warn_about_nonfunctional_selection(gs2_diagnostics_config%write_zonal_transfer, "write_zonal_transfer")

    ! Copy out internal values into module level parameters
    append_old = gs2_diagnostics_config%append_old
    conv_max_step = gs2_diagnostics_config%conv_max_step
    conv_min_step = gs2_diagnostics_config%conv_min_step
    conv_nstep_av = gs2_diagnostics_config%conv_nstep_av
    conv_nsteps_converged = gs2_diagnostics_config%conv_nsteps_converged
    conv_test_multiplier = gs2_diagnostics_config%conv_test_multiplier
    dump_check1 = gs2_diagnostics_config%dump_check1
    dump_check2 = gs2_diagnostics_config%dump_check2
    dump_fields_periodically = gs2_diagnostics_config%dump_fields_periodically
    exit_when_converged = gs2_diagnostics_config%exit_when_converged
    file_safety_check = gs2_diagnostics_config%file_safety_check
    igomega = gs2_diagnostics_config%igomega
    make_movie = gs2_diagnostics_config%make_movie
    navg = gs2_diagnostics_config%navg
    nc_sync_freq = gs2_diagnostics_config%nc_sync_freq
    nmovie = gs2_diagnostics_config%nmovie
    nsave = gs2_diagnostics_config%nsave
    nwrite = gs2_diagnostics_config%nwrite
    nwrite_mult = gs2_diagnostics_config%nwrite_mult
    ob_midplane = gs2_diagnostics_config%ob_midplane
    omegatinst = gs2_diagnostics_config%omegatinst
    omegatol = gs2_diagnostics_config%omegatol
    print_flux_line = gs2_diagnostics_config%print_flux_line
    print_line = gs2_diagnostics_config%print_line
    save_distfn = gs2_diagnostics_config%save_distfn
    save_for_restart = gs2_diagnostics_config%save_for_restart
    save_glo_info_and_grids = gs2_diagnostics_config%save_glo_info_and_grids
    save_velocities = gs2_diagnostics_config%save_velocities
    save_many = gs2_diagnostics_config%save_many
    use_nonlin_convergence = gs2_diagnostics_config%use_nonlin_convergence
    write_apar_over_time = gs2_diagnostics_config%write_apar_over_time
    write_ascii = gs2_diagnostics_config%write_ascii
    write_avg_moments = gs2_diagnostics_config%write_avg_moments
    write_bpar_over_time = gs2_diagnostics_config%write_bpar_over_time
    write_cerr = gs2_diagnostics_config%write_cerr
    write_correlation = gs2_diagnostics_config%write_correlation
    write_correlation_extend = gs2_diagnostics_config%write_correlation_extend
    write_cross_phase = gs2_diagnostics_config%write_cross_phase
    write_eigenfunc = gs2_diagnostics_config%write_eigenfunc
    write_fields = gs2_diagnostics_config%write_fields
    write_final_antot = gs2_diagnostics_config%write_final_antot
    write_final_db = gs2_diagnostics_config%write_final_db
    write_final_epar = gs2_diagnostics_config%write_final_epar
    write_final_fields = gs2_diagnostics_config%write_final_fields
    write_final_moments = gs2_diagnostics_config%write_final_moments
    write_flux_line = gs2_diagnostics_config%write_flux_line
    write_full_moments_notgc = gs2_diagnostics_config%write_full_moments_notgc
    write_g = gs2_diagnostics_config%write_g
    write_gs = gs2_diagnostics_config%write_gs
    write_gyx = gs2_diagnostics_config%write_gyx
    write_heating = gs2_diagnostics_config%write_heating
    write_kpar = gs2_diagnostics_config%write_kpar
    write_line = gs2_diagnostics_config%write_line
    write_lorentzian = gs2_diagnostics_config%write_lorentzian
    write_max_verr = gs2_diagnostics_config%write_max_verr
    write_moments = gs2_diagnostics_config%write_moments
    write_fluxes = gs2_diagnostics_config%write_fluxes
    write_fluxes_by_mode = gs2_diagnostics_config%write_fluxes_by_mode
    write_nl_flux_dist = gs2_diagnostics_config%write_nl_flux_dist
    write_omavg = gs2_diagnostics_config%write_omavg
    write_omega = gs2_diagnostics_config%write_omega
    write_parity = gs2_diagnostics_config%write_parity
    write_pflux_sym = gs2_diagnostics_config%write_pflux_sym
    write_pflux_tormom = gs2_diagnostics_config%write_pflux_tormom
    write_phi_over_time = gs2_diagnostics_config%write_phi_over_time
    write_ql_metric = gs2_diagnostics_config%write_ql_metric
    write_symmetry = gs2_diagnostics_config%write_symmetry
    write_verr = gs2_diagnostics_config%write_verr

    exist = gs2_diagnostics_config%exist

    !CMR, 12/8/2014: 
    ! Ensure write_full_moments_notgc=.false. if (write_moments .and. ob_midplane)
    ! to avoid a conflict.  
    ! FIXME: These two diagnostics do almost the same thing. Do we actually want both?
    if (write_moments .and. ob_midplane) write_full_moments_notgc=.false.

    !Override flags
    if (write_max_verr) write_verr = .true.

    print_summary = (is_list_run .and. (print_line .or. print_flux_line))

    if (is_list_run) then
       print_line = .false.
       print_flux_line = .false.
    end if

    if (no_driver .and. write_lorentzian) then
      write(*, "(a)") "WARNING: 'write_lorentzian = .true.' but antenna not enabled. Turning off 'write_lorentzian'"
      write_lorentzian = .false.
    end if

    !These don't store any data if we don't have phi so don't bother
    !calculating it.
    if(.not. has_phi) write_symmetry = .false.
    if(.not. has_phi) write_pflux_sym = .false.
    if(.not. has_phi) write_nl_flux_dist = .false.
    if(.not. has_phi) write_correlation = .false.
    if(.not. has_phi) write_correlation_extend = .false.

    if (.not. save_for_restart) nsave = -1
    write_avg_moments = write_avg_moments

    if (write_heating .and. .not. heating) then
       if (proc0) write(*,'("Warning: Disabling write_heating as collisions:heating is false.")')
       write_heating = .false.
    else if (heating .and. .not. write_heating) then
       heating = .false.
    end if

    write_any = write_line .or. write_omega     .or. write_omavg &
         .or. write_flux_line .or. write_fluxes .or. write_fluxes_by_mode  &
         .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
    dump_any = dump_check1  .or. dump_fields_periodically &
         .or.  dump_check2 .or. make_movie .or. print_summary &
         .or. write_full_moments_notgc
  end subroutine read_parameters

  !> Finalise the diagnostics module, writing final timestep diagnostics and
  !> closing any open files
  subroutine finish_gs2_diagnostics
    use mp, only: proc0
    use fields_arrays, only: phinew, bparnew
    use dist_fn, only: write_f, write_fyx, collision_error
    use gs2_io, only: nc_finish
    use antenna, only: dump_ant_amp
    use kt_grids, only: naky, ntheta0
    use gs2_time, only: user_time
    use le_grids, only: finish_weights
    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
    implicit none
    complex, dimension (ntheta0, naky) :: phi0
    logical, parameter :: last = .true.

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

    if (write_gyx) call write_fyx (phinew,bparnew,last)
    if (write_g) call write_f (last)
    if (write_cerr) call collision_error (phinew, bparnew, last)
    if (write_verr .and. proc0) call finish_weights

    ! Close some of the open ascii output files
    call close_files

    if (proc0) then
       if (write_eigenfunc) call do_write_eigenfunc(get_netcdf_file_id(), write_ascii, phi0)
       if (write_final_fields) then
          call do_write_final_fields(write_ascii, write_netcdf=.true., file_id=get_netcdf_file_id())
       end if
       if (write_kpar) call do_write_kpar(write_ascii)
       if (write_final_epar) then
          call do_write_final_epar(write_ascii, write_netcdf=.true., file_id=get_netcdf_file_id())
       end if

       ! 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(phi0, use_normalisation=write_eigenfunc, &
                                                         write_text=write_ascii, write_netcdf=.true.,&
                                                         file_id=get_netcdf_file_id())

    if (write_final_antot) then
       call do_write_final_antot(write_ascii, write_netcdf=.true., file_id=get_netcdf_file_id())
    end if

    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)

    !Finalise the netcdf file
    call nc_finish

    if (proc0) call dump_ant_amp

    if (write_gs) call do_write_gs(write_ascii)

    if (proc0) call do_write_geom

    !Now tidy up
    call deallocate_arrays

    wtmp_old = 0. ; nout = 1 ; nout_movie = 1 ; nout_big = 1
    initialized = .false.
    call gs2_diagnostics_config%reset()
  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
    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()
    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(pflux_tormom)) deallocate (pflux_tormom) 

    if (allocated(pflux_avg)) deallocate (pflux_avg, qflux_avg, heat_avg, vflux_avg)
    if (allocated(omega)) deallocate(omega, omegaavg)

    if (proc0 .and. allocated(conv_heat)) deallocate (conv_heat)

    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)
       if (write_max_verr) call close_output_file (res_unit2)
    end if
    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)
    if (dump_check1) call close_output_file (dump_check1_unit)
    if (dump_check2) call close_output_file (dump_check2_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 build_config, only: gs2_has_lowflow
    use species, only: nspec, 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 dist_fn, only: collision_error, write_f, write_fyx
    use collisions, only: ncheck, vary_vnew
    use mp, only: proc0, broadcast
    use gs2_time, only: user_time
    use gs2_io, only: nc_qflux, nc_vflux, nc_pflux, nc_pflux_tormom, nc_exchange, nc_final_fields, nc_sync, get_netcdf_movie_file_id
    use nonlinear_terms, only: nonlin
    use antenna, only: antenna_w
    use parameter_scan_arrays, only: scan_hflux => hflux_tot, scan_momflux => momflux_tot 
    use parameter_scan_arrays, only: scan_phi2_tot => phi2_tot, scan_nout => nout
    use parameter_scan, only: scanning
    use optionals, only: get_option_with_default
    use diagnostics_printout, only: write_phi_fluxes_to_unit, write_fluxes_to_unit
    use diagnostics_fluxes, only: common_calculate_fluxes, qheat, qmheat, qbheat, &
         pflux, vflux, vflux_par, vflux_perp, pflux_tormom, vflux0, vflux1, pmflux, vmflux, &
         pbflux, vbflux, exchange
    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
    real :: phi2, apar2, bpar2
    real :: t
    integer :: ik, it, is, write_mod
    complex, save :: wtmp_new !This shouldn't need to be given the save property

    real, dimension (ntheta0, nspec) :: x_qmflux
    real, dimension (nspec) ::  heat_fluxes,  part_fluxes, mom_fluxes, parmom_fluxes, perpmom_fluxes, part_tormom_fluxes
    real, dimension (nspec) :: lfmom_fluxes, vflux1_avg  ! low-flow correction to turbulent momentum fluxes
    real, dimension (nspec) :: mheat_fluxes, mpart_fluxes, mmom_fluxes
    real, dimension (nspec) :: bheat_fluxes, bpart_fluxes, bmom_fluxes
    real, dimension (nspec) :: energy_exchange
    real, dimension (nspec) ::  heat_par,  heat_perp
    real, dimension (nspec) :: mheat_par, mheat_perp
    real, dimension (nspec) :: bheat_par, bheat_perp
    real :: hflux_tot, zflux_tot, vflux_tot

    real, save :: t_old = 0.
    logical, parameter :: last = .false.
    logical:: debug=.false.
    logical :: do_force
    integer, save :: istep_last = -1
    if (present(debopt)) debug = debopt

    !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, h, hk)

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

    if (write_gyx .and. mod(istep,nmovie) == 0) call write_fyx (phinew,bparnew,last)

    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 write_f (last)

    ! 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.
    !###########################

    !Zero various arrays which may or may not get filled
    part_fluxes = 0.0 ; mpart_fluxes = 0.0 ; bpart_fluxes = 0.0
    heat_fluxes = 0.0 ; mheat_fluxes = 0.0 ; bheat_fluxes = 0.0
    mom_fluxes = 0.0 ; mmom_fluxes = 0.0 ; bmom_fluxes = 0.0  
    part_tormom_fluxes = 0.0
    energy_exchange = 0.0

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

    if (write_any_fluxes) then
       call common_calculate_fluxes()
       if (proc0) then
          if (has_phi) then
             do is = 1, nspec
                call get_volume_average (qheat(:,:,is,1), heat_fluxes(is))
                call get_volume_average (qheat(:,:,is,2), heat_par(is))
                call get_volume_average (qheat(:,:,is,3), heat_perp(is))
                call get_volume_average (pflux(:,:,is), part_fluxes(is))
                call get_volume_average (pflux_tormom(:,:,is), part_tormom_fluxes(is))
                call get_volume_average (vflux(:,:,is), mom_fluxes(is))
                call get_volume_average (vflux_par(:,:,is), parmom_fluxes(is))
                call get_volume_average (vflux_perp(:,:,is), perpmom_fluxes(is))
                call get_volume_average (exchange(:,:,is), energy_exchange(is))
#ifdef LOWFLOW
                call get_volume_average (vflux0(:,:,is), lfmom_fluxes(is))
                call get_volume_average (vflux1(:,:,is), vflux1_avg(is))
#endif
             end do
          end if
          if (has_apar) then
             do is = 1, nspec
                call get_volume_average (qmheat(:,:,is,1), mheat_fluxes(is))
                call get_surf_average (qmheat(:,:,is,1), x_qmflux(:,is))
                call get_volume_average (qmheat(:,:,is,2), mheat_par(is))
                call get_volume_average (qmheat(:,:,is,3), mheat_perp(is))
                call get_volume_average (pmflux(:,:,is), mpart_fluxes(is))
                call get_volume_average (vmflux(:,:,is), mmom_fluxes(is))
             end do
          end if
          if (has_bpar) then
             do is = 1, nspec
                call get_volume_average (qbheat(:,:,is,1), bheat_fluxes(is))
                call get_volume_average (qbheat(:,:,is,2), bheat_par(is))
                call get_volume_average (qbheat(:,:,is,3), bheat_perp(is))
                call get_volume_average (pbflux(:,:,is), bpart_fluxes(is))
                call get_volume_average (vbflux(:,:,is), bmom_fluxes(is))
             end do
          end if
          pflux_avg = pflux_avg + (part_fluxes + mpart_fluxes + bpart_fluxes)*(t-t_old)
          qflux_avg = qflux_avg + (heat_fluxes + mheat_fluxes + bheat_fluxes)*(t-t_old)
          vflux_avg = vflux_avg + (mom_fluxes + mmom_fluxes + bmom_fluxes)*(t-t_old)
          if (write_heating) heat_avg = heat_avg + h%imp_colls*(t-t_old)
          !          t_old = t
       end if

       !These are for trinity -- can we detect if this is a trinity run
       !and avoid these broadcasts when not trinity?
       call broadcast (pflux_avg)
       call broadcast (qflux_avg)
       call broadcast (vflux_avg)

       if (write_heating) call broadcast (heat_avg)
    end if

    if (proc0 .and. print_flux_line) then
      if (has_phi) then
        call write_phi_fluxes_to_unit(output_unit, t, phi2, heat_fluxes, energy_exchange)
      end if
      if (has_apar) then
        call write_fluxes_to_unit(output_unit, t, apar2, "apar", mheat_fluxes)
      end if
      if (has_bpar) then
        call write_fluxes_to_unit(output_unit, t, bpar2, "bpar", bheat_fluxes)
      end if
    end if

    ! Check for convergence
    if (nonlin .and. use_nonlin_convergence) call check_nonlin_convergence(istep, heat_fluxes(1), exit)

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

    if (.not. (write_any .or. dump_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

       hflux_tot = 0.
       zflux_tot = 0.
       vflux_tot = 0.
       if (has_phi) then
         hflux_tot = sum(heat_fluxes)
         vflux_tot = sum(mom_fluxes)
         zflux_tot = sum(part_fluxes*spec%z)
       end if
       if (has_apar) then
         hflux_tot = hflux_tot + sum(mheat_fluxes)
         vflux_tot = vflux_tot + sum(mmom_fluxes)
         zflux_tot = zflux_tot + sum(mpart_fluxes*spec%z)
       end if
       if (has_bpar) then
         hflux_tot = hflux_tot + sum(bheat_fluxes)
         vflux_tot = vflux_tot + sum(bmom_fluxes)
         zflux_tot = zflux_tot + sum(bpart_fluxes*spec%z)
       end if

       if (write_flux_line .and. write_ascii) then
         if (has_phi) then
           call write_phi_fluxes_to_unit(out_unit, t, phi2, &
                heat_fluxes, energy_exchange, part_fluxes, mom_fluxes, &
                lfmom_fluxes, vflux1_avg, gs2_has_lowflow)
         end if
         if (write_lorentzian) then
           wtmp_new = antenna_w()
           if (real(wtmp_old) /= 0. .and. 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) then
           call write_fluxes_to_unit(out_unit, t, apar2, "apar", mheat_fluxes, mpart_fluxes)
         end if
         if (has_bpar) then
           call write_fluxes_to_unit(out_unit, t, bpar2, "bpar", bheat_fluxes, bpart_fluxes)
         end if
         write (unit=out_unit, fmt="('t= ',e17.10,' h_tot= ',e11.4, ' z_tot= ',e11.4)") t, hflux_tot, zflux_tot
       end if

       if (write_fluxes) then
         call nc_qflux (get_netcdf_file_id(), nout, qheat(:,:,:,1), qmheat(:,:,:,1), qbheat(:,:,:,1), &
              heat_par, mheat_par, bheat_par, &
              heat_perp, mheat_perp, bheat_perp, &
              heat_fluxes, mheat_fluxes, bheat_fluxes, x_qmflux, hflux_tot)
         call nc_exchange (get_netcdf_file_id(), nout, exchange, energy_exchange)
         ! Update the target array in parameter_scan_arrays
         !Do you really want the scan parameter to rely on a diagnostic flag?
         scan_hflux(mod(nout-1,nstep/nwrite+1)+1) = hflux_tot

         call nc_vflux (get_netcdf_file_id(), nout, vflux, vmflux, vbflux, &
              mom_fluxes, mmom_fluxes, bmom_fluxes, vflux_tot, &
              vflux_par, vflux_perp, vflux0, vflux1)
         ! Update the target array in parameter_scan_arrays
         !Do you really want the scan parameter to rely on a diagnostic flag?
         scan_momflux(mod(nout-1,nstep/nwrite+1)+1) = vflux_tot
         call nc_pflux (get_netcdf_file_id(), nout, pflux, pmflux, pbflux, &
              part_fluxes, mpart_fluxes, bpart_fluxes, zflux_tot)
       end if

       if (write_pflux_tormom) call nc_pflux_tormom (get_netcdf_file_id(), nout, pflux_tormom, part_tormom_fluxes)

       scan_phi2_tot(mod(nout-1,nstep/nwrite+1)+1) = phi2

       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 (scanning) call bcast_scan_parameter(scan_hflux, scan_momflux, scan_phi2_tot)

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

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

    ! In new diagnostics, these two are both called for `write_symmetry`
    if (write_symmetry) call do_write_symmetry(get_netcdf_file_id(), nout)
    if (write_pflux_sym) call do_write_pflux_sym(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_avg_moments) call do_write_avg_moments

    if (write_full_moments_notgc) call do_write_full_moments_notgc

    if (dump_check1) call do_write_dump_1(t)
    if (dump_check2) call do_write_dump_2(t)

    if (dump_fields_periodically .and. mod(istep,10*nwrite) == 0) call do_dump_fields_periodically(t)

    ! Update the counter in parameter_scan_arrays
    scan_nout = nout

    !Now sync the data to file (note doesn't actually sync every call)
    if (proc0) call nc_sync(get_netcdf_file_id(), nout, get_netcdf_movie_file_id(), nout_movie, 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
  !> [[run_parameters_knobs: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 run_parameters, 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 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

    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

    if(proc0) 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)

  end subroutine do_write_ncloop

  !> Broadcast one of the input parameters according to [[parameter_scan:target_parameter_switch]]
  subroutine bcast_scan_parameter(scan_hflux, scan_momflux, scan_phi2)
    use mp, only: broadcast
    use parameter_scan, only: target_parameter_switch,target_parameter_hflux_tot
    use parameter_scan, only: target_parameter_momflux_tot,target_parameter_phi2_tot
    implicit none
    real, dimension(:), intent(in out) :: scan_hflux, scan_momflux, scan_phi2

    select case(target_parameter_switch)
    case(target_parameter_hflux_tot)
       call broadcast(scan_hflux) !This is only set if write_fluxes
    case(target_parameter_momflux_tot)
       call broadcast(scan_momflux) !This is only set if write_fluxes
    case(target_parameter_phi2_tot)
       call broadcast(scan_phi2)
    case default
       !Nothing as should generate warning/error within parameter_scan
    endselect

  end subroutine bcast_scan_parameter

  !> 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, phi0)
    use mp, only: proc0
    use file_utils, only: open_output_file, close_output_file
    use fields_arrays, only: phi, apar, bpar
    use kt_grids, only: ntheta0, naky, theta0, aky
    use theta_grid, only: theta, ntgrid
    use gs2_io, only: nc_eigenfunc
    use dist_fn, only: def_parity, even
    use run_parameters, only: has_apar
    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
    !> The normalising factor for the fields that is actually used. See
    !> [[gs2_diagnostics_knobs:write_eigenfunc]] for more details
    complex, dimension (:,:), intent(inout) :: phi0
    integer :: it, ik, ig, unit

    if(.not.proc0) return

    !Should this actually use igomega instead of 0?
    !What if fphi==0? --> See where statements below
    phi0 = phi(0,:,:)

    !This looks like a hack for the case where we know we've forced phi(theta=0) to be 0
    !this could probably be better addressed by the use of igomega above
    if (def_parity .and. has_apar .and. (.not. even)) phi0 = apar(0, :, :)

    !Address locations where phi0=0 by using next point
    where (abs(phi0) < 10.0*epsilon(0.0)) 
       phi0 = phi(1,:,:)/(theta(1)-theta(0))
    end where

    !Check again if any locations are 0, this could be true if fphi (or fapar)
    !is zero.
    where (abs(phi0) < 10.0*epsilon(0.0)) 
       phi0 = 1.0
    end where

    if (write_text) then
       call open_output_file (unit, ".eigenfunc")
       do ik = 1, naky
          do it = 1, ntheta0
             do ig = -ntgrid, ntgrid
                write (unit, "(9(1x,e12.5))") &
                     theta(ig), theta0(it,ik), aky(ik), &
                     phi(ig,it,ik)/phi0(it,ik), &
                     apar(ig,it,ik)/phi0(it,ik), &
                     bpar(ig,it,ik)/phi0(it,ik)
             end do
             write (unit, "()")
          end do
       end do
       call close_output_file (unit)
    end if

    call nc_eigenfunc (file_id, phi0)
  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_movie_file_id
    use mp, only: proc0
    implicit none
    !> Current simulation time
    real, intent(in) :: time

    if (proc0) then
       call nc_loop_movie(get_netcdf_movie_file_id(), nout_movie, time)
    end if
    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
    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
    if (.not. write_ascii) return
    allocate (j_ext(ntheta0, naky)); j_ext=0.
    call calc_jext(istep, j_ext)
    do ik=1,naky
       do it = 1, ntheta0
          if (j_ext(it,ik) .ne. 0.) 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 gs2_io, only: nc_write_moments
    use theta_grid, only: ntgrid
    use kt_grids, only: naky, ntheta0
    use species, only: nspec
    use dist_fn, only: getmoms
    use mp, only: proc0
    use fields_arrays, only: phinew, bparnew
    implicit none
    complex, dimension (-ntgrid:ntgrid,ntheta0,naky,nspec) :: ntot, density, &
         upar, tpar, tperp, qparflux, pperpj1, qpperpj1

    call getmoms (phinew, bparnew, ntot, density, upar, tpar, tperp, qparflux, pperpj1, qpperpj1)
    if(proc0) call nc_write_moments(get_netcdf_file_id(), nout, ntot, density, upar, tpar, tperp,qparflux, &
         pperpj1, qpperpj1, ob_midplane=ob_midplane)
  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
    implicit none
    !> Current simulation time
    real, intent(in) :: time
    real :: phase_tot, phase_theta
    if (.not. write_ascii) return
    call get_cross_phase (phase_tot, phase_theta)
    if (proc0) write (unit=phase_unit, fmt="('t= ',e17.10,' phase_tot= ',e11.4,' phase_theta= ',e11.4)") &
         & time, phase_tot, phase_theta
  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 dist_fn, only: get_verr, get_gtran
    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
    implicit none
    real, dimension (:,:), allocatable :: errest
    integer, dimension (:,:), allocatable :: erridx
    real :: geavg, glavg, gtavg

    allocate(errest(5,2), erridx(5,3))
    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
       ! write error estimates to .nc file          
       !          call nc_loop_vres (nout, errest_by_mode, lpcoef_by_mode)

       ! 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
          if (write_max_verr) then
             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 if
    deallocate(errest,erridx)
  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 dist_fn, 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

  !> 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 dist_fn, 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 convert, only: c2r
    use kt_grids, only: naky, ntheta0
    use fields_arrays, only: phinew, bparnew
    use dist_fn_arrays, only: gnew, g_adjust, to_g_gs2, from_g_gs2
    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, mom_fluxes_perp_dist, heat_fluxes_dist
    real, dimension (:,:), allocatable :: phi_dist
    real, dimension (:,:,:,:), allocatable :: pflux_dist, vflux_par_dist, vflux_perp_dist, qflux_dist
    real, dimension (:,:,:,:), allocatable :: phi_byMode_dist
    integer :: ig, is

    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))
    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))
    allocate (phi_byMode_dist(2,-ntgrid:ntgrid,ntheta0,naky))
    call c2r(phinew,phi_byMode_dist)
    call g_adjust (gnew, phinew, bparnew, direction = from_g_gs2) ! convert from g to h
    call flux_dist (phinew, pflux_dist, vflux_par_dist, vflux_perp_dist, qflux_dist)
    call g_adjust (gnew, phinew, bparnew, direction = to_g_gs2) ! convert back from h to g
    if (proc0) then
       if (has_phi) then
          do ig = -ntgrid, ntgrid
             do is = 1, nspec
                pflux_dist(ig,:,:,is) = pflux_dist(ig,:,:,is) * spec(is)%dens
                call get_volume_average (pflux_dist(ig,:,:,is), part_fluxes_dist(ig,is))
                vflux_par_dist(ig,:,:,is) = vflux_par_dist(ig,:,:,is) * spec(is)%dens*sqrt(spec(is)%mass*spec(is)%temp)
                call get_volume_average (vflux_par_dist(ig,:,:,is), mom_fluxes_par_dist(ig,is))
                vflux_perp_dist(ig,:,:,is) = vflux_perp_dist(ig,:,:,is) * spec(is)%dens*sqrt(spec(is)%mass*spec(is)%temp)
                call get_volume_average (vflux_perp_dist(ig,:,:,is), mom_fluxes_perp_dist(ig,is))
                qflux_dist(ig,:,:,is) = qflux_dist(ig,:,:,is) * spec(is)%dens*spec(is)%temp
                call get_volume_average (qflux_dist(ig,:,:,is), heat_fluxes_dist(ig,is))
             end do
             call get_volume_average (phi_byMode_dist(1,ig,:,:), phi_dist(1,ig))
             call get_volume_average (phi_byMode_dist(2,ig,:,:), phi_dist(2,ig))
          end do
       end if
    end if
    if (proc0) 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)
    deallocate (mom_fluxes_par_dist)
    deallocate (mom_fluxes_perp_dist)
    deallocate (heat_fluxes_dist)
    deallocate (phi_dist)
    deallocate (pflux_dist)
    deallocate (vflux_par_dist)
    deallocate (vflux_perp_dist)
    deallocate (qflux_dist)
    deallocate (phi_byMode_dist)
  end subroutine do_write_nl_flux_dist

  !> Calculate the particle flux contribution to toroidal momentum flux and write to netCDF
  subroutine do_write_pflux_sym(file_id, nout)
    use dist_fn, only: pflux_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_partsym_tormom
    implicit none
    !> NetCDF ID of the file to write to
    integer, intent(in) :: file_id
    !> Current timestep
    integer, intent(in) :: nout

    real, dimension(:,:,:), allocatable :: pflux_sym

    allocate (pflux_sym(-ntgrid:ntgrid,nlambda*negrid,nspec))
    call pflux_vs_theta_vs_vpa (pflux_sym)
    if (proc0) call nc_loop_partsym_tormom (file_id, nout, pflux_sym)
    deallocate (pflux_sym)
  end subroutine do_write_pflux_sym

  !> 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_out, 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

    if (.not. allocated(phicorr_sum)) then
       ntg_extend = (2*ntgrid+1)*((ntheta0-1)/jtwist_out+1)
       nth0_extend = min(ntheta0,jtwist_out*(naky-1))
       ! 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)
    phicorr_sum = phicorr_sum + phi_corr*(time-time_old)
    phiextend_sum = phiextend_sum + phi2_extend*(time-time_old)
    if (proc0) then
       call nc_loop_corr_extend (file_id, nout_big, time, phicorr_sum/(time-tcorr0), phiextend_sum/(time-tcorr0))
       nout_big = nout_big + 1
    end if
    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
    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
    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
    logical, save :: first = .true.

    if (.not. write_text) return

    ! initialize layouts for parity diagnostic
    if (first) then
       call init_parity_layouts (naky, nlambda, negrid, nspec, nproc, iproc)
       first = .false.
    end if

    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)
    do isgn = 1, 2
       do ig = -ntgrid, ntgrid
          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(ig,it,isgn,iplo) = gnew(ig,isgn,iglo)*aj0(ig,iglo)
                else
                   ! otherwise, send g for this iplo value to correct processor
                   call send (gnew(ig,isgn,iglo)*aj0(ig,iglo),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(ig,it,isgn,iplo),proc_id(g_lo,iglo))
             end if
          end do
       end do
    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
          do isgn = 1, 2
             sgn2 = mod(isgn,2)+1
             do ig = -ntgrid, ntgrid
                g0(ig,isgn,:) = gparity(-ig,1,sgn2,:)
             end do
          end do
       else
          do isgn = 1, 2
             sgn2 = mod(isgn,2)+1
             do ig = -ntgrid, ntgrid
                g0(ig,isgn,:) = gparity(-ig,ntheta0-it+2,sgn2,:)
             end do
          end do
       end if
       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
       do is = 1, nspec
          do ik = 1, naky
             call get_fldline_avg (gmint(:,ik,is),gmavg(it,ik,is))
             call get_fldline_avg (gpint(:,ik,is),gpavg(it,ik,is))
             call get_fldline_avg (gmnormint(:,ik,is),gmnormavg(it,ik,is))
          end do
       end do

       ! phim(theta,kx) = phi(-theta,-kx)
       ! have to treat kx=0 specially
       if (it == 1) then
          do ig = -ntgrid, ntgrid
             phim(ig,:) = phinew(-ig,1,:)
          end do
       else
          do ig = -ntgrid, ntgrid
             phim(ig,:) = phinew(-ig,ntheta0-it+2,:)
          end do
       end if

       ! 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
       do isgn = 1, 2
          do ig = -ntgrid, ntgrid
             call integrate_kysum (gm(ig,isgn,p_lo%llim_proc:p_lo%ulim_proc),ig,g_kint(ig,isgn,:),1)
          end do
       end do
       do is = 1, nspec
          call get_fldline_avg (g_kint(:,1,is)+g_kint(:,2,is),g_avg(it,is))
       end do

       ! 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
       do isgn = 1, 2
          do ig = -ntgrid, ntgrid
             call integrate_kysum (gm(ig,isgn,:),ig,gm_kint(ig,isgn,:),1)
             call integrate_kysum (gp(ig,isgn,:),ig,gp_kint(ig,isgn,:),1)
          end do
       end do
       do is = 1, nspec
          call get_fldline_avg (gm_kint(:,1,is)+gm_kint(:,2,is) &
               + gp_kint(:,1,is)+gp_kint(:,2,is),gnorm_avg(it,is))
       end do
    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
       do is = 1, nspec
          do ik = 1, naky
             call get_fldline_avg (gmint(:,ik,is),gmavg(it,ik,is))
             call get_fldline_avg (gpint(:,ik,is),gpavg(it,ik,is))
             call get_fldline_avg (gmnormint(:,ik,is),gmnormavg(it,ik,is))
          end do
       end do

       ! 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)
    do isgn = 1, 2
       do ig = -ntgrid, ntgrid
          call integrate_kysum (gp(ig,isgn,:),ig,g_kint(ig,isgn,:),1)
       end do
    end do
    do is = 1, nspec
       call get_fldline_avg (g_kint(:,1,is)+g_kint(:,2,is), g_nokx_tot(is))
    end do

    ! get normalizing terms for above diagnostic
    gm = sum(gmx,2)
    deallocate (gmx)
    do isgn = 1, 2
       do ig = -ntgrid, ntgrid
          call integrate_kysum (gm(ig,isgn,:),ig,gm_kint(ig,isgn,:),1)
          call integrate_kysum (gm(ig,isgn,:)+gp(ig,isgn,:),ig,gp_kint(ig,isgn,:),1)
       end do
    end do
    do is = 1, nspec
       call get_fldline_avg (gm_kint(:,1,is)+gm_kint(:,2,is) &
            + gp_kint(:,1,is)+gp_kint(:,2,is), gnorm_nokx_tot(is))
    end do

    ! 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
          do ig = -ntgrid, ntgrid
             g0(ig,:,:) = gparity(-ig,1,:,:)
          end do
       else
          do ig = -ntgrid, ntgrid
             g0(ig,:,:) = gparity(-ig,ntheta0-it+2,:,:)
          end do
       end if
       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
       do is = 1, nspec
          do ik = 1, naky
             call get_fldline_avg (gmint(:,ik,is),gmavg(it,ik,is))
             call get_fldline_avg (gpint(:,ik,is),gpavg(it,ik,is))
             call get_fldline_avg (gmnormint(:,ik,is),gmnormavg(it,ik,is))
          end do
       end do

       ! phim(theta,kx) = phi(-theta,-kx)
       ! have to treat kx=0 specially
       if (it == 1) then
          do ig = -ntgrid, ntgrid
             phim(ig,:) = phinew(-ig,1,:)
          end do
       else
          do ig = -ntgrid, ntgrid
             phim(ig,:) = phinew(-ig,ntheta0-it+2,:)
          end do
       end if

       ! 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
       do ig = -ntgrid, ntgrid
          call integrate_kysum (gm(ig,1,:)+gm(ig,2,:),ig,g_kint(ig,1,:),1)
       end do
       do is = 1, nspec
          call get_fldline_avg (g_kint(:,1,is),g_avg(it,is))
       end do

       ! 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
       do ig = -ntgrid, ntgrid
          call integrate_kysum (gm(ig,1,:)+gm(ig,2,:),ig,gm_kint(ig,1,:),1)
          call integrate_kysum (gp(ig,1,:)+gp(ig,2,:),ig,gp_kint(ig,1,:),1)
       end do
       do is = 1, nspec
          call get_fldline_avg (gm_kint(:,1,is)+gp_kint(:,1,is),gnorm_avg(it,is))
       end do
    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
       do ik = 1, naky
          do it = 1, ntheta0
             call get_fldline_avg (real(gparity(:,it,ik,is)),gmavg(it,ik,is))
          end do
       end do
       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

  !> Calculate flux surfgace averaged low-order moments of \(g\) and write to netCDF
  subroutine do_write_avg_moments
    use dist_fn, only: getmoms
    use mp, only: proc0
    use kt_grids, only: ntheta0, naky
    use species, only: nspec
    use fields_arrays, only: phinew, bparnew
    use gs2_io, only: nc_loop_moments
    use theta_grid, only: ntgrid
    implicit none
    integer :: it, is
    complex, dimension (-ntgrid:ntgrid,ntheta0,naky,nspec) :: ntot, density, &
         upar, tpar, tperp, qparflux, pperpj1, qpperpj1
    complex, dimension (ntheta0, nspec) :: ntot00, density00, upar00, tpar00, tperp00
    complex, dimension (ntheta0) :: phi00
    real, dimension (nspec) :: ntot2, ntot20, tpar2, tperp2
    real, dimension (ntheta0, naky, nspec) :: ntot2_by_mode, ntot20_by_mode
    real, dimension (ntheta0, naky, nspec) :: tpar2_by_mode, tperp2_by_mode

    ! We seem to call this routine a lot in one loop
    call getmoms (phinew, bparnew, ntot, density, upar, tpar, tperp, qparflux, pperpj1, qpperpj1)

    if (proc0) then
       do is = 1, nspec
          do it = 1, ntheta0
             call get_fldline_avg(ntot(-ntgrid:ntgrid,it,1,is),ntot00(it,is))
             call get_fldline_avg(density(-ntgrid:ntgrid,it,1,is),density00(it,is))
             call get_fldline_avg(upar(-ntgrid:ntgrid,it,1,is),upar00(it,is))
             call get_fldline_avg(tpar(-ntgrid:ntgrid,it,1,is),tpar00(it,is))
             call get_fldline_avg(tperp(-ntgrid:ntgrid,it,1,is),tperp00(it,is))
          end do
       end do

       do it = 1, ntheta0
          call get_fldline_avg(phinew(-ntgrid:ntgrid,it,1),phi00(it))
       end do

       do is = 1, nspec
          call get_vol_average (ntot(:,:,:,is), ntot(:,:,:,is), &
               ntot2(is), ntot2_by_mode(:,:,is))
          call get_vol_average (tpar(:,:,:,is), tpar(:,:,:,is), &
               tpar2(is), tpar2_by_mode(:,:,is))
          call get_vol_average (tperp(:,:,:,is), tperp(:,:,:,is), &
               tperp2(is), tperp2_by_mode(:,:,is))
       end do

       do is = 1, nspec
          call get_vol_average (ntot(0,:,:,is), ntot(0,:,:,is), &
               ntot20(is), ntot20_by_mode(:,:,is))
       end do

       call nc_loop_moments (get_netcdf_file_id(), nout, ntot2, ntot2_by_mode, ntot20, &
            ntot20_by_mode, phi00, ntot00, density00, upar00, &
            tpar00, tperp00, tpar2_by_mode, tperp2_by_mode)

    end if
  end subroutine do_write_avg_moments

  !> Calculate moments (density, parallel flow, and parallel and perpendicular
  !> temperatures) in non-guiding centre coordinates and write to netCDF
  subroutine do_write_full_moments_notgc()
    use dist_fn, only: getmoms_notgc
    use theta_grid, only: ntgrid
    use kt_grids, only: naky, ntheta0
    use species, only: nspec
    use mp, only: proc0
    use gs2_io, only: nc_loop_fullmom
    use fields_arrays, only: phinew, bparnew
    implicit none
    complex, dimension (-ntgrid:ntgrid,ntheta0,naky,nspec) :: ntot, density, &
         upar, tpar, tperp
    call getmoms_notgc(phinew, bparnew, density,upar,tpar,tperp,ntot)
    if(proc0) then
       call nc_loop_fullmom(get_netcdf_file_id(), nout, &
            & ntot(igomega,:,:,:),density(igomega,:,:,:), &
            & upar(igomega,:,:,:), &
            & tpar(igomega,:,:,:),tperp(igomega,:,:,:) )
    endif
  end subroutine do_write_full_moments_notgc

  !> Calculate the field-line average of \(\phi\) and write to [[dump_check1_unit]]
  subroutine do_write_dump_1(t)
    use mp, only: proc0
    use kt_grids, only: naky, ntheta0, aky, akx
    use theta_grid, only: field_line_average
    use fields_arrays, only: phinew
    use dist_fn, only: omega0, gamma0
    use constants, only: zi
    implicit none
    real, intent(in) :: t
    complex :: tmp, sourcefac
    integer :: ik, it

    if(.not.proc0) return

    !Should we not actually use the sourcefac calculated in dist_fn for consistency?
    sourcefac = exp(-zi*omega0*t+gamma0*t)

    do ik = 1, naky
       do it = 1, ntheta0
          tmp=field_line_average(phinew(:,it,ik))
          write (dump_check1_unit, "(20(1x,e12.5))") &
               t, aky(ik), akx(it), tmp, tmp/sourcefac
       end do
    end do
  end subroutine do_write_dump_1

  !> Calculate \(A_\parallel(k_x, k_y)\) at `igomega` and write to [[dump_check2_unit]]
  subroutine do_write_dump_2(t)
    use mp, only: proc0
    use kt_grids, only: naky, ntheta0, aky, akx
    use fields_arrays, only: aparnew
    implicit none
    real, intent(in) :: t
    integer :: it, ik

    if(.not.proc0) return
    do ik = 1, naky
       do it = 1, ntheta0
          write (dump_check2_unit, "(5(1x,e12.5))") &
               t, aky(ik), akx(it), aparnew(igomega,it,ik)
       end do
    end do
  end subroutine do_write_dump_2

  !> Write out \(\phi, A_\parallel, B_\parallel\) to text file
  subroutine do_dump_fields_periodically(t)
    use file_utils, only: open_output_file, close_output_file
    use kt_grids, only: naky, ntheta0, aky, theta0, akx
    use theta_grid, only: theta, ntgrid
    use fields_arrays, only: phinew, aparnew, bparnew
    use constants, only: run_name_size
    use mp, only: proc0
    implicit none
    real, intent(in) :: t
    character(run_name_size) :: filename
    integer :: ik, it, ig, unit

    if (.not. proc0) return

    write (filename, "('dump.fields.t=',e13.6)") t
    call open_output_file(unit, '.'//filename)
    do ik = 1, naky
       do it = 1, ntheta0
          do ig = -ntgrid, ntgrid
             write (unit=unit, fmt="(20(1x,e12.5))") &
                  theta(ig), aky(ik), theta0(it,ik), &
                  phinew(ig,it,ik), aparnew(ig,it,ik), &
                  bparnew(ig,it,ik), t, akx(it)
          end do
          write (unit, "()")
       end do
    end do
    call close_output_file (unit)
  end subroutine do_dump_fields_periodically

  !> 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_parity) flush (parity_unit)
  end subroutine flush_files

  !> Trinity 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 gs2_time, only: user_time
    use mp, only: proc0, broadcast, job
    logical, intent(inout) :: exit
    integer, intent(in) :: istep
    real, intent(in) :: heat_flux

    real :: heat_av_new, heat_av_diff
    integer :: place, iwrite, nwrite_av
    logical :: debug = .false.

    if(proc0 .and. .not. (trin_istep .ne. 0 .and. istep .eq. 0)) then
       if(istep .gt. 0) trin_istep = trin_istep + nwrite ! Total number of steps including restarted trinity runs
       iwrite = trin_istep/nwrite ! Number of diagnostic write steps written
       nwrite_av = conv_nstep_av/nwrite ! Number of diagnostic write steps to average over        

       heat_sum_av = (heat_sum_av * iwrite + heat_flux) / (iwrite+1) ! Cumulative average of heat flux

       place = mod(trin_istep,conv_nstep_av)/nwrite
       conv_heat(place) = heat_flux

       if (debug) write(6,'(A,I5,A,e11.4,A,I6,I6)') 'Job ',job, &
            ' time = ',user_time, ' step = ',trin_istep
       if (debug) write(6,'(A,I5,A,e11.4,A,e11.4)') 'Job ',job, &
            ' heat = ',heat_flux, ' heatsumav = ',heat_av

       if (trin_istep .ge. conv_nstep_av) then
          heat_av_new = sum(conv_heat) / nwrite_av
          heat_av_diff = heat_av_new - heat_av
          if(debug) write(6,'(A,I5,A,e11.4,A,e11.4)') 'Job ',job, &
               ' heat_sum_av_diff = ',heat_sum_av
          heat_av = heat_av_new
          ! Convergence test - needs to be met conv_nsteps_converged/nwrite times in succession
          if (abs(heat_av_diff) .lt. heat_av_test) then
             conv_isteps_converged = conv_isteps_converged + 1
          else
             conv_isteps_converged = 0
             heat_av_test = heat_sum_av * conv_test_multiplier
          endif

          if ((conv_isteps_converged .ge. conv_nsteps_converged/nwrite) .and. &
               (trin_istep .ge. conv_min_step)) then
             if (debug) write(6,'(A,I5,A,I6,I3)')'Job ',job, &
                  ' Reached convergence condition after step ',trin_istep
             exit = .true. .and. exit_when_converged
          endif

          if(trin_istep .gt. conv_max_step) then
             write(6,'(A,I5,A,I7)') '*** Warning. Job ',job, &
                  ' did not meet the convergence condition after ',trin_istep
             exit = .true. .and. exit_when_converged
          endif
       endif
    endif

    call broadcast(exit)

  end subroutine check_nonlin_convergence

  !> Calculate some heating diagnostics, and update their time history
  subroutine heating (istep, h, hk)
    use mp, only: proc0
    use dist_fn, only: get_heat
    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 dist_fn_arrays, only: c_rate
    use gs2_heating, only: heating_diagnostics, avg_h, avg_hk, zero_htype
    use collisions, only: collisions_heating => heating
    implicit none
    integer, intent (in) :: istep
    !> Heating diagnostics summed over space at the current timestep
    type (heating_diagnostics), intent(in out) :: h
    !> Heating diagnostics as a function of \((\theta, k_y)\) at the current timestep
    type (heating_diagnostics), dimension(:,:), intent(in out) :: hk

    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

    call get_jext(j_ext)

    !Do averages with a history variable
    if (proc0) then
       !Save variable to history
       if (navg > 1) then
          ! 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 if
    end if

  end subroutine calc_jext

  !> A linear estimate of the diffusivity, used for Trinity testing
  function diffusivity()
    use kt_grids, only: ntheta0, naky, kperp2
    use theta_grid, only: grho
    real :: diffusivity
    real, dimension(ntheta0, naky) :: diffusivity_by_k
    integer  :: ik, it
    diffusivity_by_k = 0.0
    do ik=1,naky
       do it=1,ntheta0
          if (kperp2(igomega,it,ik).eq.0.0) cycle
          diffusivity_by_k(it,ik) = &
               max(aimag(omegaavg(it,ik)), 0.0)/kperp2(igomega, it, ik)*2.0
       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
  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
    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 (aky(ik) == 0) 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
    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 (tol_use == 0.0) 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
    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=.false.

    if (debug) write(6,*) "get_omeaavg: allocate domega"
    if (.not. allocated(domega)) allocate(domega(navg,ntheta0,naky))
    if (present(debopt)) debug=debopt
    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.eq.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) 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')")
          exit = .true. .and. exit_when_converged
       end if

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

  !> Summed magnitude of all the fields
  subroutine phinorm (phitot)
    use fields_arrays, only: phinew, aparnew, bparnew
    use theta_grid, only: theta
    use kt_grids, only: naky, ntheta0
    use constants, only: pi
    use integration, only: trapezoidal_integration
    implicit none
    real, dimension (:,:), intent (out) :: phitot
    integer :: ik, it
    do ik = 1, naky
       do it = 1, ntheta0
          phitot(it,ik) = 0.5/pi &
               *(trapezoidal_integration(theta, &
               (abs(phinew(:,it,ik))**2 + abs(aparnew(:,it,ik))**2 &
               + abs(bparnew(:,it,ik))**2)))
       end do
    end do
  end subroutine phinorm

  !> FIXME : Add documentation
  !!
  !! @note Here we calculate the correction factors for each possible kperp
  subroutine get_vol_average_all (a, b, axb, axb_by_mode)
    use theta_grid, only: ntgrid, field_line_average
    use kt_grids, only: naky, ntheta0
    implicit none
    complex, dimension (-ntgrid:,:,:), intent (in) :: a, b
    real, intent (out) :: axb
    real, dimension (:,:), intent (out) :: axb_by_mode
    integer :: ik, it
    integer :: ng

    ng = ntgrid

    do ik = 1, naky
       do it = 1, ntheta0
          axb_by_mode(it,ik) &
               = field_line_average(real(conjg(a(-ng:ng,it,ik))*b(-ng:ng,it,ik)))
       end do
    end do

    call get_volume_average (axb_by_mode, axb)
  end subroutine get_vol_average_all

  !> FIXME : Add documentation  
  subroutine get_vol_average_one (a, b, axb, axb_by_mode)
    implicit none
    complex, dimension (:,:), intent (in) :: a, b
    real, intent (out) :: axb
    real, dimension (:,:), intent (out) :: axb_by_mode

    axb_by_mode = real(conjg(a)*b)

    call get_volume_average (axb_by_mode, axb)

  end subroutine get_vol_average_one

  !> FIXME : Add documentation  
  subroutine get_volume_average (f, favg)
    use kt_grids, only: naky, ntheta0, aky
    implicit none
    real, dimension (:,:), intent (in) :: f
    real, intent (out) :: favg
    real :: fac
    integer :: ik, it

    ! ky=0 modes have correct amplitudes; rest must be scaled
    ! note contrast with scaling factors in FFT routines.

    !CMR+GC: 2/9/2013
    !  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)
    favg = 0.
    do ik = 1, naky
       fac = 0.5
       if (aky(ik) == 0.) fac = 1.0
       do it = 1, ntheta0
          favg = favg + f(it, ik) * fac
       end do
    end do

  end subroutine get_volume_average

  !> FIXME : Add documentation  
  subroutine get_surf_average (f, favg)
    use kt_grids, only: naky, ntheta0, aky
    implicit none
    real, dimension (:,:), intent (in) :: f
    real, dimension (:), intent (out) :: favg
    real :: fac
    integer :: ik, it

    ! ky=0 modes have correct amplitudes; rest must be scaled
    ! note contrast with scaling factors in FFT routines.

    !CMR+GC: 2/9/2013
    !  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)

    favg = 0.
    do ik = 1, naky
       fac = 0.5
       if (aky(ik) == 0.) fac = 1.0
       do it = 1, ntheta0
          favg(it) = favg(it) + f(it, ik) * fac
       end do
    end do

  end subroutine get_surf_average

  !> FIXME : Add documentation  
  subroutine reset_init

    use gs2_time, only: user_time
    use mp, only: proc0

    implicit none

    ! HJL > reset values for convergence condition
    if(proc0) then
       conv_isteps_converged = 0
       heat_sum_av = 0.0
       conv_heat = 0.0
       trin_istep = 0
    endif
    start_time = user_time
    pflux_avg = 0.0 ; qflux_avg = 0.0 ; heat_avg = 0.0 ; vflux_avg = 0.0

  end subroutine reset_init

  !> FIXME : Add documentation  
  subroutine ensemble_average (nensembles, time_int)
    use mp, only: scope, allprocs, group_to_all, broadcast
    use gs2_time, only: user_time
    use species, only: nspec
    implicit none
    integer, intent (in) :: nensembles
    real, intent (out) :: time_int
    integer :: is
    real, dimension (nensembles) :: dt_global
    real, dimension (nensembles,nspec) :: pflx_global, qflx_global, heat_global, vflx_global
    time_int=user_time-start_time
    call scope (allprocs)
    call group_to_all (time_int, dt_global, nensembles)
    call broadcast (dt_global)
    time_int = sum(dt_global)
    call group_to_all (pflux_avg, pflx_global, nensembles)
    call group_to_all (qflux_avg, qflx_global, nensembles)
    call group_to_all (vflux_avg, vflx_global, nensembles)

    call broadcast (pflx_global)
    call broadcast (qflx_global)
    call broadcast (vflx_global)
    do is = 1, nspec
       pflux_avg = sum(pflx_global(:,is))
       qflux_avg = sum(qflx_global(:,is))
       vflux_avg = sum(vflx_global(:,is))
    end do
    if (write_heating) then
       call group_to_all (heat_avg, heat_global, nensembles)
       call broadcast (heat_global)
       do is = 1, nspec
          heat_avg = sum(heat_global(:,is))
       end do
    end if
  end subroutine ensemble_average

  !> Calculate the field-line average (real values)
  subroutine get_fldline_avg_r (fld_in, fld_out)
    use theta_grid, only: field_line_average, ntgrid
    implicit none
    real, dimension (-ntgrid:), intent (in) :: fld_in
    real, intent (out) :: fld_out
    fld_out = field_line_average(fld_in)
  end subroutine get_fldline_avg_r

  !> Calculate the field-line average (complex values)
  subroutine get_fldline_avg_c (fld_in, fld_out)
    use theta_grid, only: field_line_average, ntgrid
    implicit none
    complex, dimension (-ntgrid:), intent (in) :: fld_in
    complex, intent (out) :: fld_out
    fld_out = field_line_average(fld_in)
  end subroutine get_fldline_avg_c

  !> This is a highly simplified synthetic diagnostic which 
  !! calculates the cross phase between the electron density and the 
  !! perpendicular electron temperature for comparisons with DIII-D.  
  !! Returns the value of the cross-phase at the outboard midplane and 
  !! integrated over all v and x. We can generalize this routine to 
  !! other fields at some point, but for now this is just a skeleton for 
  !! a more realistic synthetic diagnostic.
  subroutine get_cross_phase (phase_tot, phase_theta)
    use species, only: nspec, spec, is_electron_species
    use kt_grids, only: ntheta0, naky
    use theta_grid, only: ntgrid
    use dist_fn, only: getemoms
    use fields_arrays, only: phinew, bparnew
    implicit none
    real, intent (out) :: phase_tot, phase_theta
    complex, dimension (:,:,:,:), allocatable :: ntot, tperp
    complex, dimension (ntheta0, naky) :: nTp_by_mode
    complex :: nTp
    !    real, dimension (ntheta0, naky) :: n2_by_mode, T2_by_mode
    !    real :: n2, T2

    integer ::  is

    allocate ( ntot(-ntgrid:ntgrid,ntheta0,naky,nspec))
    allocate (tperp(-ntgrid:ntgrid,ntheta0,naky,nspec))

    call getemoms (phinew, bparnew, ntot, tperp)

    do is = 1,nspec
       if (is_electron_species(spec(is))) then
          ! get cross_phase at outboard midplane
          call get_vol_int (ntot(0,:,:,is), tperp(0,:,:,is), nTp, nTp_by_mode)
          phase_theta = atan2(aimag(nTp),real(nTp))!/sqrt(n2*T2)
          ! get integrated cross_phase 
          call get_vol_int (ntot(:,:,:,is), tperp(:,:,:,is), nTp, nTp_by_mode)
          phase_tot = atan2(aimag(nTp),real(nTp))!/sqrt(n2*T2)
       end if
    end do

    deallocate (ntot, tperp)

  end subroutine get_cross_phase

  !> FIXME : Add documentation
  subroutine get_vol_int_all (a, b, axb, axb_by_mode)
    use theta_grid, only: ntgrid, field_line_average
    use kt_grids, only: naky, ntheta0
    implicit none
    complex, dimension (-ntgrid:,:,:), intent (in) :: a, b
    complex, intent (out) :: axb
    complex, dimension (:,:), intent (out) :: axb_by_mode
    integer :: ik, it
    integer :: ng
    ng = ntgrid

    do ik = 1, naky
       do it = 1, ntheta0
          axb_by_mode(it,ik) &
               = field_line_average(conjg(a(-ng:ng,it,ik))*b(-ng:ng,it,ik))
       end do
    end do

    call get_volume_int (axb_by_mode, axb)
  end subroutine get_vol_int_all

  !> FIXME : Add documentation  
  subroutine get_vol_int_one (a, b, axb, axb_by_mode)
    implicit none
    complex, dimension (:,:), intent (in) :: a, b
    complex, intent (out) :: axb
    complex, dimension (:,:), intent (out) :: axb_by_mode

    axb_by_mode = conjg(a)*b
    call get_volume_int (axb_by_mode, axb)

  end subroutine get_vol_int_one

  !> FIXME : Add documentation    
  subroutine get_volume_int (f, favg)
    use kt_grids, only: naky, ntheta0, aky
    implicit none
    complex, dimension (:,:), intent (in) :: f
    complex, intent (out) :: favg
    real :: fac
    integer :: ik, it

    ! ky=0 modes have correct amplitudes; rest must be scaled
    ! note contrast with scaling factors in FFT routines.

    !CMR+GC: 2/9/2013
    !  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)

    favg = 0.
    do ik = 1, naky
       fac = 0.5
       if (aky(ik) == 0.) fac = 1.0
       do it = 1, ntheta0
          favg = favg + f(it, ik) * fac
       end do
    end do

  end subroutine get_volume_int

  !> Calculate the correlation function over the extended domain
  subroutine correlation_extend (cfnc, phi2extend)
    !    use constants, only: pi
    use fields_arrays, only: phinew
    use theta_grid, only: ntgrid, jacob, delthet, ntgrid
    !   use theta_grid, only: theta
    use kt_grids, only: ntheta0, naky, jtwist_out

    implicit none

    real, dimension (:,:,:), intent (out) :: phi2extend
    complex, dimension (:,:,:), intent (out) :: cfnc

    integer :: ig, it, ik, im, igmod
    integer :: itshift, nconnect, offset
    !    real :: fac

    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_out)*(2*ntgrid+1)/2
    do ig = -ntgrid, ntgrid
       call reorder_kx (phinew(ig,:,1), phir(ig,:))
    end do
    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
       do ig = -ntgrid, ntgrid
          call reorder_kx (phinew(ig,:,ik), phir(ig,:))
       end do
       ! shift in kx due to parallel boundary condition
       ! also the number of independent theta0s
       itshift = jtwist_out*(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_out-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

  !> 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

  !> FIXME : Add documentation    
  subroutine reorder_kx (unord, ord)
    use kt_grids, only: ntheta0
    implicit none
    complex, dimension (:), intent (in) :: unord
    complex, dimension (:), intent (out) :: ord

    ord(:ntheta0/2) = unord(ntheta0/2+2:)
    ord(ntheta0/2+1:) = unord(:ntheta0/2+1)

  end subroutine reorder_kx

  !> Set the module level config type
  !> Will abort if the module has already been initialised to avoid
  !> inconsistencies.
  subroutine set_gs2_diagnostics_config(gs2_diagnostics_config_in)
    use mp, only: mp_abort
    type(gs2_diagnostics_config_type), intent(in), optional :: gs2_diagnostics_config_in
    if (initialized) then
       call mp_abort("Trying to set gs2_diagnostics_config when already initialized.", to_screen = .true.)
    end if
    if (present(gs2_diagnostics_config_in)) then
       gs2_diagnostics_config = gs2_diagnostics_config_in
    end if
  end subroutine set_gs2_diagnostics_config

  !> Get the module level config instance
  function get_gs2_diagnostics_config()
    type(gs2_diagnostics_config_type) :: get_gs2_diagnostics_config
    get_gs2_diagnostics_config = gs2_diagnostics_config
  end function get_gs2_diagnostics_config
  
end module gs2_diagnostics