diagnostics_zonal_m.f90 Source File


Contents


Source Code

module diagnostics_zonal_m
    !! Module for computing zonal diagnostics for the BRAGINSKII model
    use netcdf
    use MPI
    use comm_handler_m, only : comm_handler_t
    use screen_io_m, only :  get_stdout
    use diagnostic_variable_m, only : diagnostic_variable_t
    use variable_m, only : variable_t
    use precision_grillix_m, only : GP, MPI_GP
    use perf_m, only : perf_start, perf_stop
    use error_handling_grillix_m, only: handle_error_netcdf, handle_error, error_info_t
    use status_codes_grillix_m, only : GRILLIX_ERR_OTHER
    use sources_external_m, only : sources_external_t
    use mesh_cart_m, only: mesh_cart_t
    use parallel_map_m, only : parallel_map_t
    use inplane_operators_m, only : inplane_operators_t
    use equilibrium_m, only: equilibrium_t
    use penalisation_m, only : penalisation_t
    use polars_m, only : polars_t
    use equilibrium_storage_m, only : equilibrium_storage_t
    use parallel_target_flux_m, only : parallel_target_flux_t
    use multistep_m, only : multistep_storage_t
    use polar_operators_m, only : flux_surface_average, surface_average, &
                                  drift_surface_integral

    use diagnostics_helpers_m, only: inner_product, total_volume
    use diagnostics_group_m, only : diagnostics_group_t
    use diagnostics_packet_m, only : diagnostics_packet_t
    use gyroviscosity_m, only : eta_neo_flow

    use params_brag_model_m, only : rhos, beta

    implicit none

    type, public, extends(diagnostics_group_t) :: &
                                diagnostics_zonal_t
    contains
        procedure, public :: process_to_file
        procedure, public :: project_zonal
    end type

contains

    subroutine process_to_file(self, diag_packet, comm_handler, equi, equi_on_cano, equi_on_stag, mesh_cano, &
                            mesh_stag, map, penalisation_cano, penalisation_stag, polars_cano, polars_stag, &
                            tau, ne, te, ti, pot, vort, upar, jpar, apar, apar_fluct, &
                            neutrals_dens, neutrals_parmom, neutrals_pressure, &
                            src_ne, src_te, src_ti, src_upar, src_vort, &
                            isnaps, idiag, start_new_file)
        !! Main diagnostics computation routine
        class(diagnostics_zonal_t), intent(inout) :: self
        !! Instance of class
        type(diagnostics_packet_t), intent(inout) :: diag_packet
        !! Diagnostic packet
        type(comm_handler_t), intent(in) :: comm_handler
        !! MPI communication handler
        class(equilibrium_t), intent(in) :: equi
        !! Magnetic equilibrium
        class(equilibrium_storage_t), intent(in) :: equi_on_cano
        !! Equilibrium storage on canonical plane enabling faster performance at certain locations
        class(equilibrium_storage_t), intent(in) :: equi_on_stag
        !! Equilibrium storage on staggered plane enabling faster performance at certain locations
        type(mesh_cart_t), intent(in) :: mesh_cano
        !! Mesh (canonical) within poloidal plane
        type(mesh_cart_t), intent(in) :: mesh_stag
        !! Mesh (staggered) within poloidal plane
        type(parallel_map_t), intent(in) :: map
        !! Parallel map 
        type(penalisation_t), intent(in) :: penalisation_cano
        !! Penalisation (canonical)
        type(penalisation_t), intent(in) :: penalisation_stag
        !! Penalisation (staggered)
        type(polars_t), intent(in) :: polars_cano
        !! Polar grid and operators (canonical)
        type(polars_t), intent(in) :: polars_stag
        !! Polar grid and operators (staggered)
        real(GP), intent(in) :: tau
        !! Time
        type(variable_t), intent(in) :: ne 
        !! Electron density
        type(variable_t), intent(in) :: te
        !! Electron temperature
        type(variable_t), intent(in) :: ti
        !! Ion temperature
        type(variable_t), intent(in) :: pot
        !! Electrostatic potential
        type(variable_t), intent(in) :: vort
        !! Generalised vorticity
        type(variable_t), intent(in) :: upar
        !! Parallel ion velocity
        type(variable_t), intent(in) :: jpar
        !! Parallel current
        type(variable_t), intent(in) :: apar
        !! Parallel electromagnetic potential
        type(variable_t), intent(in) :: apar_fluct
        !! Parallel electromagnetic potential fluctation
        type(variable_t), intent(in) :: neutrals_dens
        !! Neutrals density
        type(variable_t), intent(in) :: neutrals_parmom
        !! Neutrals parallel momentum
        type(variable_t), intent(in) :: neutrals_pressure
        !! Neutrals pressure
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: src_ne
        !! Particle source values on inner mesh points
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: src_te
        !! Electron temperature source values on inner mesh points
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: src_ti
        !! Ion temperature source values on inner mesh points
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: src_upar
        !! Parallel momentum source values on inner mesh points
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: src_vort
        !! Vorticity source values on inner mesh points
        integer, intent(in) :: isnaps
        !! Braginskii model snapshot number
        integer, intent(in) :: idiag
        !! Diagnostic snapshot number
        logical, intent(in) :: start_new_file
        !! Flag whether to write to new file
        
        integer :: ind
        !! Reference index of last initialized diagnostic
        integer :: ip, nrho, l, ki
        real(GP), dimension(mesh_cano%get_n_points()) :: wrk, aux
        real(GP), dimension(polars_cano%grid%get_nrho()) :: tmp_zon

        nrho = polars_cano%grid%get_nrho()

        call perf_start('diagnostics_braginskii_zonal')

        call self%allocate_diags(n_diags=86) ! << ADJUST WHEN ADDING DIAGNOSTICS
    
        ! Time
        call self%init_next_diagnostic(1, 'dim_scalar', 'tau', ind)
        self%diags(ind)%vals(1) = tau

        ! Polar coordinates
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'rho', ind)
        do ip = 1, nrho
            self%diags(ind)%vals(ip) = polars_cano%grid%get_rho(ip)
        enddo

        ! Zonal volume averages of base variables
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ne_zonavg', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                ne%vals, mode='ZONAVG')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'te_zonavg', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                te%vals, mode='ZONAVG')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ti_zonavg', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                ti%vals, mode='ZONAVG')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'pot_zonavg', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                pot%vals, mode='ZONAVG')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'upar_zonavg', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                upar%vals, mode='ZONAVG')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'neutrals_dens_zonavg', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                neutrals_dens%vals, mode='ZONAVG')

        ! Surface integrals of radial gradients
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'radgrad_ne_srfint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%radgrad_ne, mode='SRFINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'radgrad_te_srfint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%radgrad_te, mode='SRFINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'radgrad_ti_srfint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%radgrad_ti, mode='SRFINT')

        ! Zonal volume integrals of external sources
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_ne_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%src_ext_ne, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_te_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%src_ext_te, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_ti_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%src_ext_ti, mode='ZONINT')

        ! Zonal volume integrals of changerates of base variables
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_ne_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ddt_ne, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_te_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ddt_te, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_ti_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ddt_ti, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_neutrals_dens_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ddt_neutrals_dens, mode='ZONINT')

        ! Plasma-neutrals sources
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_iz_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%src_iz_ne, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_rec_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%src_rec_ne, mode='ZONINT')
        ! call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_rcy_zonint', ind)
        ! call self%project_zonal(comm_handler, equi_on_cano, mesh_cano, map, polars_cano, ind, &
        !                         diag_packet%src_rcy, mode='ZONINT')

        ! Numerical sources
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'src_floor_ne_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%src_floor_ne, mode='ZONINT')


        ! -----------------------------------------------------------------------------------------
        ! Changerate of energetic quantities
        ! -----------------------------------------------------------------------------------------

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_ddt_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_ddt, mode='ZONINT')

         call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_eth_zonint', ind)
         call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                 diag_packet%ddt_eth, mode='ZONINT')

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_ekin_perp_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ekin_perp_ddt, mode='ZONINT')

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_ekin_par_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ekin_par_ddt, mode='ZONINT')

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ddt_eem_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%eem_ddt, mode='ZONINT')

        ! -----------------------------------------------------------------------------------------
        ! Particle fluxes
        ! -----------------------------------------------------------------------------------------

        ! ExB particle flux
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'cont_exb_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), pot%vals, self%diags(ind)%vals, ne%vals)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Diamagnetic particle flux
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, ne, te, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -ne%vals(l) * te%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'cont_dia_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), wrk, self%diags(ind)%vals)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Electromagnetic flutter particle flux
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, ne, upar, jpar, wrk)
        !$omp do
            do l = 1, mesh_cano%get_n_points()
                wrk(l) = -(diag_packet%ne_stag(l) * upar%vals(l) - jpar%vals(l)) * equi_on_cano%btor(l)
            enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'cont_flutter_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! -----------------------------------------------------------------------------------------
        ! Thermal electron energy fluxes
        ! -----------------------------------------------------------------------------------------
        
        ! ExB flux of electron thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, ne, te, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = 1.5_GP * ne%vals(l) * te%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_exb_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), pot%vals, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Diamagnetic flux of electron thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, ne, te, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -2.5_GP * ne%vals(l) * te%vals(l)**2
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_dia_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), wrk, self%diags(ind)%vals)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Flutter flux of advection of electron thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, ne, upar, jpar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -2.5_GP * (diag_packet%ne_stag(l) * upar%vals(l) - jpar%vals(l)) &
                     * diag_packet%te_stag(l) * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_flutter_adv_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Flutter flux of heat flux of electron thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, ne, upar, jpar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -diag_packet%qpar_cond_e(l) * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_flutter_heat_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Flutter flux of thermal force heat flux of electron thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, jpar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = 0.71_GP * diag_packet%te_stag(l) * jpar%vals(l) * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_flutter_thermalforceheat_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! -----------------------------------------------------------------------------------------
        ! Thermal ion energy fluxes
        ! -----------------------------------------------------------------------------------------

        ! ExB flux of ion thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = 1.5_GP * ne%vals(l) * ti%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_exb_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), pot%vals, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Diamagnetic flux of ion thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = 2.5_GP * ne%vals(l) * ti%vals(l)**2
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_dia_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), wrk, self%diags(ind)%vals)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Flutter flux of advection of ion thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, ne, upar, jpar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -2.5_GP * diag_packet%ne_stag(l) * upar%vals(l) * diag_packet%ti_stag(l) &
                     * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_flutter_adv_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Flutter flux of heat flux of ion thermal energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, ne, upar, jpar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -diag_packet%qpar_cond_i(l) * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_flutter_heat_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Polarisation flux, as volume integral
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_upol_flux_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethi_upol_flux, mode='ZONINT')

        ! -----------------------------------------------------------------------------------------
        ! Perpendicular kinetic energy fluxes
        ! -----------------------------------------------------------------------------------------

        ! ExB flux of perpendicular kinetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, mesh_cano, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = diag_packet%ekin_perp(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_exb_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), pot%vals, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Diamagnetic flux of perpendicular kinetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, mesh_cano, ne, ti, wrk, aux)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) =  diag_packet%ekin_perp(l) / ne%vals(l)
            aux(l) = ne%vals(l) * ti%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_dia_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), aux, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Polarisation flux, as volume integral
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_upol_flux_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ekinperp_upol_flux, mode='ZONINT')

        ! Flutter flux of perpendicular kinetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, wrk, aux)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -diag_packet%ekin_perp(l) * diag_packet%upar_cano(l) * equi_on_cano%btor(l)
            aux(l) = diag_packet%aparv_fluct_cano(l) / equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_flutter_adv_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), aux, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Diamagnetic and ExB part of Poynting flux
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, mesh_cano, ne, ti, te, pot, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) =  pot%vals(l) * ne%vals(l) * (ti%vals(l) + te%vals(l)) 
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_poynt_diaexb_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), wrk, self%diags(ind)%vals)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Polarisation part of Poynting flux
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_jpol_poynt_flux_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ekinperp_jpol_poynt_flux, mode='ZONINT')

        ! Flutter part of Poynting flux
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, equi_on_cano, diag_packet, jpar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -jpar%vals(l) * diag_packet%pot_stag(l) * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_poynt_flutter_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Gyroviscous perpendicular flux
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_gvisc_flux', ind)

        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, diag_packet, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -diag_packet%gstressv(l) / 3.0_GP
        enddo
        !$omp end do
        !$omp end parallel
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), pot%vals, tmp_zon, wrk)

        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, diag_packet, ne, ti, aux, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -diag_packet%gstressv(l) / (3.0_GP * ne%vals(l))
            aux(l) = ne%vals(l) * ti%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), aux, self%diags(ind)%vals, wrk)

        self%diags(ind)%vals = (self%diags(ind)%vals + tmp_zon) * rhos

        ! -----------------------------------------------------------------------------------------
        ! Parallel kinetic energy fluxes
        ! -----------------------------------------------------------------------------------------

        ! ExB flux of parallel kinetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, mesh_cano, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = diag_packet%ekin_par(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinpar_exb_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%pot_stag, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Diamagnetic flux of parallel kinetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, mesh_cano, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) =  diag_packet%ekin_par(l) * diag_packet%ti_stag(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinpar_dia_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), wrk, self%diags(ind)%vals)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! Flutter flux of advection of parallel kinetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, equi_on_cano, mesh_cano, ne, ti, wrk, upar)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -diag_packet%ekin_par(l) * upar%vals(l) * equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinpar_flutter_adv_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), diag_packet%apar_fluct_over_btor, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta

        ! Flutter flux of gyroviscousity
        !$omp parallel default(none) private(l) &
        !$omp shared(diag_packet, equi_on_cano, mesh_cano, upar, wrk, aux)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = -2.0_GP/3.0_GP * diag_packet%gstressv(l) * diag_packet%upar_cano(l) * equi_on_cano%btor(l)
            aux(l) = diag_packet%aparv_fluct_cano(l) / equi_on_cano%btor(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinpar_flutter_gvisc_flux', ind)
        call drift_surface_integral(comm_handler%get_comm_planes(), equi, mesh_cano, polars_cano, &
                                    map%get_dphi(), aux, self%diags(ind)%vals, wrk)
        self%diags(ind)%vals = self%diags(ind)%vals * rhos * beta


        ! -----------------------------------------------------------------------------------------
        ! Electromagnetic energy fluxes
        ! -----------------------------------------------------------------------------------------

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'eem_flux_srfint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%eem_flux, mode='SRFINT')
        self%diags(ind)%vals = self%diags(ind)%vals * rhos

        ! -----------------------------------------------------------------------------------------
        ! Transfer terms in energy theorem
        ! -----------------------------------------------------------------------------------------

        ! -pe * div(v_{ExB})
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_-pe_div_exb', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_pe_div_exb, mode='ZONINT')

        ! -pi * div(v_{ExB})
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_-pi_div_exb', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethi_pi_div_exb, mode='ZONINT')

        ! vpar * grad_par(p_e)
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_vpar_gradpar_pe', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_vpar_gradpar_pe, mode='ZONINT')

        ! upar * grad_par(p_i)
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_upar_gradpar_pi', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethi_upar_gradpar_pi, mode='ZONINT')

        ! upar * grad_par(p_e)
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_upar_gradpar_pe', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%etrans_upar_gradpar_pe, mode='ZONINT')

        ! upol * grad(p_i)
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_upol_grad_pi', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%etrans_upol_grad_pi, mode='ZONINT')
        ! jpar grad_par(pot)
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_jpar_gradpar_pot', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%etrans_jpar_gradpar_pot, mode='ZONINT')
        
        ! Ohmic heating: eta * jpar^2
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_eta_jpar_sqrd', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_eta_jpar2, mode='ZONINT')

        ! Thermal force term: -0.71 jpar * grad_par(te)
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_-0.71_jpar_gradpar_te', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_jpar_gradpar_te, mode='ZONINT')

        ! Equipartition term: -Q_delta 
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_-qdelta', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_equipartition, mode='ZONINT')

        ! Gyroviscous heating term
        !$omp parallel default(none) private(l) shared(mesh_cano, diag_packet, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = diag_packet%gstressv(l)**2 &
                     / (3.0_GP * eta_neo_flow(ne%vals(l), ti%vals(l)))
        enddo
        !$omp end do
        !$omp end parallel                  
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_gvisc_heating', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Gyroviscous kinetic energy transfer term
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, diag_packet, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = - diag_packet%gstressv(l) * diag_packet%gstressv_par(l) &
                     / (3.0_GP * eta_neo_flow(ne%vals(l), ti%vals(l)))
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'etrans_gvisc_kin', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! -----------------------------------------------------------------------------------------
        ! Source terms (Non-Braginskii)
        ! -----------------------------------------------------------------------------------------

        ! Total particle source
        wrk = 0.0_GP
        !$omp parallel default(none) private(ki, l) shared(mesh_cano, src_ne, wrk)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            wrk(l) = src_ne(ki)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'cont_src', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')
        
        ! Electron thermal energy source
        wrk = 0.0_GP
        !$omp parallel default(none) private(ki, l) shared(mesh_cano, ne, te, src_ne, src_te, wrk)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            wrk(l) = 1.5_GP * (src_ne(ki) * te%vals(l) + src_te(ki) * ne%vals(l))
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_src', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                               wrk, mode='ZONINT')

        ! Ion thermal energy source
        wrk = 0.0_GP
        !$omp parallel default(none) private(ki, l) shared(mesh_cano, ne, ti, src_ne, src_ti, wrk)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            wrk(l) = 1.5_GP * (src_ne(ki) * ti%vals(l) + src_ti(ki) * ne%vals(l))
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_src', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Perpendicular kinetic energy source
        wrk = 0.0_GP
        !$omp parallel default(none) private(ki, l) &
        !$omp shared(mesh_cano, diag_packet, src_ne, src_vort, ne, pot, wrk)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            wrk(l) = diag_packet%ekin_perp(l) / ne%vals(l) * src_ne(ki) &
                     + pot%vals(l) * src_vort(ki)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_src', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Parallel kinetic energy source
        wrk = 0.0_GP
        !$omp parallel default(none) private(ki, l) &
        !$omp shared(mesh_cano, diag_packet, src_ne, src_upar, upar, wrk)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            ! Note: Here we add quantities on staggered and canonical mesh_cano
            ! This is fine here since the projection below is a linear operation
            wrk(l) = 0.5_GP * diag_packet%upar_cano(l)**2 * src_ne(ki) &
                     + src_upar(ki) * diag_packet%ne_stag(l) * upar%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinpar_src', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! -----------------------------------------------------------------------------------------
        ! Dissipation terms (Non-Braginskii)
        ! -----------------------------------------------------------------------------------------

        ! Particle flux due to hyperviscosity and buffer zones
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'cont_perpvisc_flux', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%cont_radflux_visc, mode='SRFINT')

        ! Electron thermal energy dissipation
        !$omp parallel default(none) private(l) shared(mesh_cano, diag_packet, ne, te, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = 1.5_GP * (diag_packet%diss_ne(l) * te%vals(l) + &
                               diag_packet%diss_te(l) * ne%vals(l) + &
                               diag_packet%diss_pe(l))
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_diss', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Viscous electron thermal flux due to hyperviscosity and buffer zones
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_perpvisc_flux', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_radflux_visc, mode='SRFINT')

        ! Ion thermal energy dissipation
        !$omp parallel default(none) private(l) shared(mesh_cano, diag_packet, ne, ti, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = 1.5_GP * (diag_packet%diss_ne(l) * ti%vals(l) + &
                               diag_packet%diss_ti(l) * ne%vals(l) + &
                               diag_packet%diss_pi(l))
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_diss', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Viscous ion thermal flux due to hyperviscosity and buffer zones
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_perpvisc_flux', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethi_radflux_visc, mode='SRFINT')

        ! Perpendicular kinetic energy dissipation
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, diag_packet, ne, pot, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = diag_packet%ekin_perp(l) / ne%vals(l) * diag_packet%diss_ne(l) &
                     - pot%vals(l) * diag_packet%diss_vort(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinperp_diss', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Parallel kinetic energy dissipation
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, diag_packet, upar, wrk)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            ! Note: Here we add quantities on staggered and canonical mesh_cano
            ! This is fine here since the projection below is a linear operation
            wrk(l) = 0.5_GP * diag_packet%upar_cano(l)**2 * diag_packet%diss_ne(l) &
                     + diag_packet%diss_upar(l) * diag_packet%ne_stag(l) * upar%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ekinpar_diss', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! Dissipation of electromagnetic energy
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, diag_packet, jpar, wrk, beta)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            wrk(l) = beta * jpar%vals(l) * diag_packet%diss_apar(l)
        enddo
        !$omp end do
        !$omp end parallel
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'eem_diss', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                wrk, mode='ZONINT')

        ! -----------------------------------------------------------------------------------------
        ! Further zonal diagnostic terms
        ! -----------------------------------------------------------------------------------------

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_hypvisc_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_hypvisc, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_buffer_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_buffer, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_src_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_src, mode='ZONINT')

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethe_radflux_mag_qnormal_srfint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethe_radflux_mag_qnormal, mode='SRFINT')

        call self%init_next_diagnostic(nrho, 'dim_zonal', 'ethi_radflux_mag_qnormal_srfint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%ethi_radflux_mag_qnormal, mode='SRFINT')

        ! Vorticity equation: volumetric terms
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'vort_curv_pi_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%vort_curv_pi, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'vort_div_jpar_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%vort_div_jpar, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'vort_flutter_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%vort_flutter, mode='ZONINT')
        call self%init_next_diagnostic(nrho, 'dim_zonal', 'vort_hypvisc_zonint', ind)
        call self%project_zonal(comm_handler, equi, equi_on_cano, mesh_cano, map, polars_cano, ind, &
                                diag_packet%vort_hypvisc, mode='ZONINT')

        call self%write_diagnostics(tau, isnaps, idiag, start_new_file)
        deallocate(self%diags)

        call perf_stop('diagnostics_braginskii_zonal')

    end subroutine

    subroutine project_zonal(self, comm_handler, equi, equi_storage, mesh, map, polars, ind, u, mode)
        !! Project input onto zonal diagnostic dimension nrho
        class(diagnostics_zonal_t), intent(inout) :: self
        !! Instance of class
        type(comm_handler_t), intent(in) :: comm_handler
        !! MPI communication handler
        class(equilibrium_t) :: equi
        class(equilibrium_storage_t), intent(in) :: equi_storage
        !! Equilibrium storage enabling faster performance at certain locations
        type(mesh_cart_t), intent(in) :: mesh
        !! Mesh within poloidal plane
        type(parallel_map_t), intent(in) :: map
        !! Parallel map 
        type(polars_t), intent(in) :: polars
        !! Polar grid and operators
        integer, intent(in) :: ind
        !! Work index in diagnostics array
        real(GP), dimension(:), intent(in) :: u
        !! Input array
        character(len=*), intent(in) :: mode
        !! Projection mode, select from : 
        !!      'ZONAVG' - zonal volume average aka 'flux surface average'
        !!      'ZONINT' - zonal volume integral
        !!      'SRFAVG' - surface average
        !!      'SRFINT' - surface integral

        integer :: ip, comm

        comm = comm_handler%get_comm_planes()

        select case(mode)
        case('ZONAVG')
            call flux_surface_average(comm, mesh, polars, u, self%diags(ind)%vals)
        case('ZONINT')
            call flux_surface_average(comm, mesh, polars, u, self%diags(ind)%vals)
            do ip = 1, polars%grid%get_nrho()
                self%diags(ind)%vals(ip) = self%diags(ind)%vals(ip) * polars%fluxsurf_vol(ip)
            enddo
        case('SRFAVG')
            call surface_average(comm, mesh, polars, u, self%diags(ind)%vals)
        case('SRFINT')
            call surface_average(comm, mesh, polars, u, self%diags(ind)%vals)
            do ip = 1, polars%grid%get_nrho()
                self%diags(ind)%vals(ip) = self%diags(ind)%vals(ip) * polars%fluxsurf_area(ip)
            enddo
        case default
             call handle_error('Invalid mode for zonal projection', &
                              GRILLIX_ERR_OTHER, __LINE__, __FILE__)
        end select
    end subroutine

end module