polarisation_braginskii_m.f90 Source File


Contents


Source Code

module polarisation_braginskii_m
    !! This module handles the polarisation part in the quasi-neutrality equation, particularly the time discretisation
    !! 
    !! \f[
    !!      \nabla\cdot\left[co\frac{\partial}{\partial t}\nabla_\perp\mathbf{E}_\perp\right] = rhs, 
    !! \f]
    !! where
    !! \f[
    !!      \mathbf{E}_\perp = \left(\nabla_\perp\phi+\tau\frac{\nabla_\perp p_i}{n}\right)
    !! \f]
    !! the perpendicular electric field.
    !! 
    !! The module solves for the generalised vorticity atr timestep t+1
    !! \f[
    !!      Omega^{t+1} = div[co  (\nabla_perp\pot^{t+1}) 
    !! \f]
    !! and electrostatic potential
    !! 
    !! The time-advancement is quite intricate due to the position of the time derivative in the vorticity equation
    !! The scheme is described in [W. Zholobenko et al., Contrib. Plasma Phys. 2019;e201900131. (https://doi.org/10.1002/ctpp.201900131.)]
    !! Therefore some intricate modifications to the storage values in the time-step integrator tstep_vort are necessary and additional storage of values
    !!
    !! Also, parallel advection of vorticity is handeled here.
    !!
    !! Further intricate things arise from the use of Shortley-Weller stencil for the elliptic operator
    !!
    use MPI
    use NETCDF
    use comm_handler_m, only : comm_handler_t
    use precision_grillix_m, only : GP, GP_NAN
    use descriptors_braginskii_m, only : BND_BRAGTYPE_ZONAL_NEUMANN
    use error_handling_grillix_m, only: handle_error_netcdf
    use multistep_m, only : multistep_storage_t, karniadakis_t
    use parallel_map_m, only : parallel_map_t
    use penalisation_m, only : penalisation_t
    use polars_m, only : polars_t
    use polar_operators_m, only : surface_average
    use boundaries_braginskii_m, only : boundaries_braginskii_t
    use variable_m, only: variable_t
    ! From PARALLAX  
    use perf_m, only : perf_start, perf_stop
    use interpolation_m, only: interpol1d
    use equilibrium_m, only : equilibrium_t
    use mesh_cart_m, only : mesh_cart_t 
    use helmholtz_solver_m, only : helmholtz_solver_t                           
    use helmholtz_operator_m, only : helmholtz_single_inner
    use boundaries_perp_m, only : set_perpbnds, extrapolate_ghost_points
    use descriptors_m, only : BND_TYPE_NONE, BND_TYPE_NEUMANN, BND_TYPE_DIRICHLET_ZERO, DISTRICT_CORE 
    use descriptors_braginskii_m, only : BND_BRAGTYPE_FLOATING_POTENTIAL_LOCAL
    use helmholtz_solver_m, only : helmholtz_solver_t
    use helmholtz_solver_factory_m, only : parameters_helmholtz_solver_factory, helmholtz_solver_factory
    use helmholtz_netcdfio_m, only : write_netcdf_helmholtz
    ! Parameters
    use params_tstep_m, only : &
        tstep_order
    use params_brag_model_m, only : &
        rhos, tratio
    use params_brag_boundaries_perp_m, only : &
        bnddescr_pot_core, bnddescr_pot_wall, bnddescr_pot_dome, bnddescr_pot_out
    use params_brag_boundaries_parpen_m, only : &
        bnddescr_pot_pen, pen_epsinv
    use params_brag_switches_m, only : &
        swb_vort_dia, swb_vort_paradv
    implicit none

    public :: polarisation_advance
    public :: compute_vorticity
    private :: pot_solve
    
 contains

    subroutine polarisation_advance(comm_handler, &
                                 equi, mesh_cano, hsolver_cano, penalisation_cano, polars_cano, &
                                 boundaries, tstep_vort, &   
                                 pot_tseq, ne_tseq, ti_tseq, pot_firstsolve_store, &
                                 co, nev, tiv, dvort, pot_pen, &
                                 pot_adv, vort_adv, &
                                 sinfo, res, sinfo_zon, res_zon)
        !! Advances vorticity in time and solves for pential
        type(comm_handler_t), intent(in) :: comm_handler
        !! Communicators
        class(equilibrium_t), intent(in) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(in) :: mesh_cano
        !! Mesh (canonical)
        class(helmholtz_solver_t), intent(inout) :: hsolver_cano
        !! Elliptic (2D) solver on canonical mesh
        type(penalisation_t), intent(in) :: penalisation_cano
        !! Penalisation (canonical)
        type(polars_t), intent(in) :: polars_cano
        !! Polar grid and operators (canonical)
        type(boundaries_braginskii_t), intent(inout) :: boundaries
        !! Boundary information for the BRAGINSKII model 
        type(karniadakis_t), intent(inout) :: tstep_vort
        !! Time-step integrator for vorticity equation
        real(GP), dimension(mesh_cano%get_n_points(), tstep_order), intent(in) :: pot_tseq
        !! Time sequence of electrostatic potential at time-points, t, t-1, t-2,...t-(order-1)
        real(GP), dimension(mesh_cano%get_n_points(), tstep_order), intent(in) :: ne_tseq
        !! Time sequence of density values at time-points, t, t-1, t-2,...t-(order-1)
        real(GP), dimension(mesh_cano%get_n_points(), tstep_order), intent(in) :: ti_tseq
        !! Time sequqence of ion temperature values at time-points, t, t-1, t-2,...t-(order-1)
        type(multistep_storage_t), intent(inout) :: pot_firstsolve_store
        !! Storage of electrostatic potential values at time-points, t, t-1, t-2,...t-(order-1) -
        !! first solve - only used with BND_BRAGTYPE_ZONAL_NEUMANN
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: co
        !! Polarisation coefficient
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: nev
        !! Density at t+1
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: tiv
        !! Ion temperature at t+1
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: dvort
        !! Right hand side of vorticity equation
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: pot_pen
        !! Penalisation values for electrostatic potential 
        real(GP), dimension(mesh_cano%get_n_points()), intent(inout) :: pot_adv
        !! Initial guess on input / values of electrostatic potential at time t+1 on output
        real(GP), dimension(mesh_cano%get_n_points()), intent(out) :: vort_adv
        !! Values of Vorticity at time t+1 
        integer, intent(out) :: sinfo
        !! Infor from solver
        real(GP), intent(out) :: res
        !! Residual of solve       
        integer, intent(out) :: sinfo_zon
        !! Info of solver required for zonal (second) solve in case of zonal Neumann boundary condition
        real(GP), intent(out) :: res_zon
        !! Residual of zonal (second) solve in case of zonal Neumann boundary condition       

        integer :: nstorage, istorage
        real(GP), dimension(mesh_cano%get_n_points()) :: wrk          
        integer :: ki, kb, l, nrho, district
        real(GP) :: pen_fac, drho, rho_core_rel, pot_core
        
        real(GP), dimension(polars_cano%grid%get_nrho()) :: pot_zon
        real(GP), dimension(mesh_cano%get_n_points_boundary()) :: bndvals_potzon
        
        integer, parameter :: intorder=3

        integer :: bnddescr_actual_pot_core, bnddescr_actual_pot_wall, bnddescr_actual_pot_dome, bnddescr_actual_pot_out

        call perf_start('../../polarisation_advance')

        call perf_start('../../../compute_vort')
        ! Compute vorticity at t-1, t-2, with co at t+1
        ! Store result in respective storage at time-step integrator 
        nstorage = tstep_vort%get_current_order()-1
        do istorage = 1, nstorage
            call compute_vorticity(equi, mesh_cano, boundaries, co, &
                                   potv=pot_tseq(:,istorage+1), &
                                   nev=ne_tseq(:,istorage+1), & 
                                   tiv=ti_tseq(:,istorage+1), &
                                   vortv=tstep_vort%vstore(:,istorage,1))
        enddo

        ! Compute vorticity at t
        call compute_vorticity(equi, mesh_cano, boundaries, co, &
                               potv=pot_tseq(:,1), &
                               nev=ne_tseq(:,1), &  
                               tiv=ti_tseq(:,1), &
                               vortv=wrk)

        ! Advance vorticity to t+1 
        call tstep_vort%advance(dvort, wrk, vort_adv)
        call tstep_vort%shift_storage((/wrk, dvort/)) 
    
        ! Penalisation (diagonal implicit part on vorticity)
        !$omp parallel default(none) private(ki, l, pen_fac) &
        !$omp          shared(mesh_cano, penalisation_cano, boundaries, tstep_vort, pen_epsinv, vort_adv)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            pen_fac = (1.0_GP + penalisation_cano%get_charfun_val(ki) * pen_epsinv * tstep_vort%get_dtau_bdf())
            vort_adv(l) = vort_adv(l) / pen_fac
        enddo
        !$omp end do

        ! Set boundary conditions on vorticity
        call set_perpbnds(mesh_cano, boundaries%vort%types, vort_adv, boundaries%vort%vals)
        call extrapolate_ghost_points(mesh_cano, vort_adv)
        !$omp end parallel

        ! Compute diamagnetic part of voticity at t+1
        call compute_vorticity(equi, mesh_cano, boundaries, &
                               co=co, &
                               potv=nev, &  ! Dummy not used as swpot = 0
                               nev=nev, tiv=tiv, vortv=wrk, swpot=0.0_GP)

        call perf_stop('../../../compute_vort')
        
        ! Set boundary condition for potential solve (basic parallax descriptors) -----------------
        if (bnddescr_pot_core == BND_BRAGTYPE_ZONAL_NEUMANN) then
            bnddescr_actual_pot_core = BND_TYPE_NEUMANN
        elseif (bnddescr_pot_core == BND_BRAGTYPE_FLOATING_POTENTIAL_LOCAL) then
            bnddescr_actual_pot_core = BND_TYPE_DIRICHLET_ZERO    
        else
            bnddescr_actual_pot_core = bnddescr_pot_core
        endif
        
        if (bnddescr_pot_wall == BND_BRAGTYPE_FLOATING_POTENTIAL_LOCAL) then
            bnddescr_actual_pot_wall = BND_TYPE_DIRICHLET_ZERO    
        else
            bnddescr_actual_pot_wall = bnddescr_pot_wall
        endif

        if (bnddescr_pot_dome == BND_BRAGTYPE_FLOATING_POTENTIAL_LOCAL) then
            bnddescr_actual_pot_dome = BND_TYPE_DIRICHLET_ZERO    
        else
            bnddescr_actual_pot_dome = bnddescr_pot_dome
        endif

        if (bnddescr_pot_out == BND_BRAGTYPE_FLOATING_POTENTIAL_LOCAL) then
            bnddescr_actual_pot_out = BND_TYPE_DIRICHLET_ZERO    
        else
            bnddescr_actual_pot_out = bnddescr_pot_out
        endif

        ! Solve for electrostatic potential -------------------------------------------------------
        call perf_start('../../../pot_solve')
        call pot_solve(comm_handler, equi, mesh_cano, hsolver_cano, penalisation_cano, & 
                       bnddescr_actual_pot_core, bnddescr_actual_pot_wall, bnddescr_actual_pot_dome, bnddescr_actual_pot_out, &
                       boundaries%pot%vals, &
                       co, wrk, pot_pen,vort_adv, pot_adv, &
                       sinfo, res)
        call perf_stop('../../../pot_solve')

        ! Re-solve for zonal Neumann boundary condition with obtained zonal average value at core as Dirichlet condition at core
        
        if (bnddescr_pot_core == BND_BRAGTYPE_ZONAL_NEUMANN) then
            call perf_start('../../../pot_solve2')
            
            bnddescr_actual_pot_core = BND_TYPE_DIRICHLET_ZERO ! Set boundary condition for potential core

            ! Store result of first solve to have a good initial guess for the next time-step
            call pot_firstsolve_store%shift_storage(pot_adv)

            ! Compute zonal value of potential at core
            nrho = polars_cano%grid%get_nrho()
            drho = polars_cano%grid%get_drho()
            call surface_average(comm_handler%get_comm_planes(), mesh_cano, polars_cano, pot_adv, pot_zon)
            rho_core_rel = equi%rhomin - polars_cano%grid%get_rhopol_min()
            pot_core = interpol1d(intorder, nrho, drho, pot_zon, rho_core_rel)

            ! Solve with pot = pot_core as dirichlet boundary condition at core
            !$omp parallel default(none) private(kb, l, district) &
            !$omp          shared(mesh_cano, pot_tseq, pot_adv, bndvals_potzon, boundaries, pot_core)
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                pot_adv(l) = 3.0_GP*pot_tseq(l,1) - 3.0_GP*pot_tseq(l,2) + pot_tseq(l,3)
            enddo
            !$omp end do
            !$omp do
            do kb = 1, mesh_cano%get_n_points_boundary()
                l = mesh_cano%boundary_indices(kb)
                district = mesh_cano%get_district(l)  
                select case(district) 
                    case(DISTRICT_CORE)
                        bndvals_potzon(kb) = pot_core
                    case default
                        bndvals_potzon(kb) = boundaries%pot%vals(kb)
                end select
            end do
            !$omp end do
            !$omp end parallel
                
            call pot_solve(comm_handler, equi, mesh_cano, hsolver_cano, penalisation_cano, &
                           bnddescr_actual_pot_core, bnddescr_actual_pot_wall, bnddescr_actual_pot_dome, bnddescr_actual_pot_out, &         
                           bndvals_potzon, &  
                           co, wrk, pot_pen, vort_adv, pot_adv, &
                           sinfo_zon, res_zon)    
            
            call perf_stop('../../../pot_solve2')

        else
            sinfo_zon = 0
            res_zon     = 0.0_GP
        endif

        call perf_stop('../../polarisation_advance')

    end subroutine

                                                                  
    subroutine compute_vorticity(equi, mesh_cano, boundaries, co, potv, nev, tiv, vortv, swpot, swdia)
        !! Computes generalised vorticity
        !! \f[
        !!      \nabla\cdot\left[co\nabla_\perp\left(swpot\nabla_\perp\phi+swdia\tau\frac{\nabla_\perp p_i}{n}\right)\right] 
        !! \f]
        class(equilibrium_t), intent(in) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(in) :: mesh_cano
        !! Mesh (canonical)
        type(boundaries_braginskii_t), intent(in) :: boundaries
        !! Boundary information for the BRAGINSKII model 
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: co
        !! Polarisation coefficient
        real(GP), dimension(mesh_cano%get_n_points()) :: potv
        !! Values of electrostatic potential
        real(GP), dimension(mesh_cano%get_n_points()) :: nev
        !! Values of density
        real(GP), dimension(mesh_cano%get_n_points()) :: tiv
        !! Values of ion temperature
        real(GP), dimension(mesh_cano%get_n_points()), intent(out) :: vortv
        !! Generalised values for vorticity on inner grid
        real(GP), intent(in), optional :: swpot
        !! Switch in front of \nabla_\perp\phi term, default = 1
        real(GP), intent(in), optional :: swdia
        !! Switch in front of \nabla_\perp p_i/n term, default = 1

        integer :: ki, kb, kg, l
        real(GP) :: fac, phi_cano, x, y, fac_pot, fac_dia, xi, ell_pot, ell_dia
        real(GP), dimension(mesh_cano%get_n_points()) :: codia, piv

        phi_cano = mesh_cano%get_phi()

        ! Set up switches
        if (present(swpot)) then
            fac_pot = swpot       
        else
            fac_pot = 1.0_GP
        endif
        if (present(swdia)) then
            fac_dia = swdia * swb_vort_dia * tratio   
        else
            fac_dia = swb_vort_dia * tratio   
        endif   
        
        ! Set up coefficients
        fac = -rhos**2
        !$omp parallel default(none) &
        !$omp          private(ki, kb, kg, l, x, y, xi, ell_pot, ell_dia) &
        !$omp          shared(equi, mesh_cano, phi_cano, fac, fac_pot, fac_dia, &
        !$omp                 co, nev, tiv,codia, piv, potv, vortv)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            codia(l) = co(l) / nev(l)
            piv(l) = nev(l) * tiv(l)
        enddo
        !$omp end do
        !$omp do
        do kb = 1, mesh_cano%get_n_points_boundary()
            l = mesh_cano%boundary_indices(kb)
            codia(l) = co(l) / nev(l)
            piv(l) = nev(l) * tiv(l)
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh_cano%get_n_points_ghost()
            l = mesh_cano%ghost_indices(kg)
            codia(l) = co(l) / nev(l)
            piv(l) = nev(l) * tiv(l)
        enddo
        !$omp end do
        ! Apply elliptic operators
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            x = mesh_cano%get_x(l)
            y = mesh_cano%get_y(l) 
            xi = fac / equi%jacobian(x, y, phi_cano)
            ell_pot = helmholtz_single_inner(mesh_cano, potv, l, co, xi) 
            ell_dia = helmholtz_single_inner(mesh_cano, piv, l, codia, xi) 

            vortv(l) = fac_pot * ell_pot + fac_dia * ell_dia
        enddo
        !$omp end do
        !$omp end parallel

        ! Set ghost points on vorticity
        call set_perpbnds(mesh_cano, boundaries%vort%types, vortv, boundaries%vort%vals)
        call extrapolate_ghost_points(mesh_cano, vortv)
       
    end subroutine


    subroutine pot_solve(comm_handler, equi, mesh_cano, hsolver_cano, penalisation_cano, &
                         bnddescr_actual_pot_core, bnddescr_actual_pot_wall, bnddescr_actual_pot_dome, bnddescr_actual_pot_out, &
                         bndvals_pot, &
                         co, diavort, pot_pen, vortv, potv, sinfo, res)
        !! Solves for electrostatic potential
        type(comm_handler_t), intent(in) :: comm_handler
        !! Communicators
        class(equilibrium_t), intent(in) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(in) :: mesh_cano
        !! Mesh (canonical)
        class(helmholtz_solver_t), intent(inout) :: hsolver_cano
        !! Elliptic (2D) splver on canonical mesh
        type(penalisation_t), intent(in) :: penalisation_cano
        !! Penalisation (canonical)
        integer, intent(in) ::  bnddescr_actual_pot_core
        !! Boundary descriptor for core 
        integer, intent(in) ::  bnddescr_actual_pot_wall
        !! Boundary descriptor for wall 
        integer, intent(in) ::  bnddescr_actual_pot_dome
        !! Boundary descriptor for dome 
        integer, intent(in) ::  bnddescr_actual_pot_out
        !! Boundary descriptor for out 
        real(GP), dimension(mesh_cano%get_n_points_boundary()), intent(in) :: bndvals_pot
        !! Boundary information for the BRAGINSKII model 
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: co
        !! Polarisation coefficient
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: diavort
        !! Diamagnetic part of generalised vorticity (only inner points accessed)
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(in) :: pot_pen
        !! Penalisation values for electrostatic potential 
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: vortv
        !! Generalised vorticty
        real(GP), dimension(mesh_cano%get_n_points()), intent(inout) :: potv
        !! Electrostatic potential
        integer, intent(out) :: sinfo
        !! Info from solver
        real(GP), intent(out) :: res
        !! Residual of solution      

        real(GP), dimension(mesh_cano%get_n_points()) :: rhs
        real(GP), dimension(mesh_cano%get_n_points_inner()) :: lambda, xi            
        integer :: ki, kb, kg, l, nf90_stat, nf90_id
        real(GP) :: phi_cano, x, y
        real(GP) :: fac, pen_char    
        character(len=3) :: plane_c
        
        fac = rhos**2
        
        phi_cano = mesh_cano%get_phi()

        ! Setup equation system (coefficients and right hand side)
        call perf_start('../../../../pot_solve_setup')
        !$omp parallel default(none) &
        !$omp          private(ki, kb, kg, l, x, y, pen_char) &
        !$omp          shared(equi, mesh_cano, phi_cano, penalisation_cano, bndvals_pot, &
        !$omp                 pen_epsinv, fac, bnddescr_pot_pen, &
        !$omp                 lambda, xi, rhs, vortv, diavort, pot_pen)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()        
            l = mesh_cano%inner_indices(ki) 
            x = mesh_cano%get_x(l)
            y = mesh_cano%get_y(l) 
            xi(ki) = fac / equi%jacobian(x, y, phi_cano)
            rhs(l) = -vortv(l) + diavort(l)
            
            ! Apply penalisation for potential solve
            if (bnddescr_pot_pen == BND_TYPE_NONE) then               
                lambda(ki) = 0.0_GP  
            else
                pen_char = penalisation_cano%get_charfun_val(ki)
                lambda(ki) = pen_char * pen_epsinv * fac
                xi(ki) = xi(ki) * (1.0_GP - pen_char) 
                rhs(l) = rhs(l) * (1.0_GP - pen_char) + pen_char * pen_epsinv * fac * pot_pen(ki)
            endif
        enddo
        !$omp end do
        !$omp do
        do kb = 1, mesh_cano%get_n_points_boundary()        
            l = mesh_cano%boundary_indices(kb) 
            rhs(l) = bndvals_pot(kb)
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh_cano%get_n_points_ghost()        
            l = mesh_cano%ghost_indices(kg) 
            rhs(l) = 0.0_GP
        enddo
        !$omp end do
        !$omp end parallel

        call perf_stop('../../../../pot_solve_setup')
        
        call perf_start('../../../../pot_solve_init')

        call hsolver_cano%update(co, lambda, xi, &
                                 bnddescr_actual_pot_core, &
                                 bnddescr_actual_pot_wall, &
                                 bnddescr_actual_pot_dome, &
                                 bnddescr_actual_pot_out)
        call perf_stop('../../../../pot_solve_init')
        
        call perf_start('../../../../pot_solve_solve')
        call hsolver_cano%solve(rhs, potv, res, sinfo)
        call perf_stop('../../../../pot_solve_solve')
        
        if (sinfo < 0) then
            ! Write out state if solver failed
            write(plane_c,'(I3.3)')comm_handler%get_plane()
            
            ! Overwrites existing file
            nf90_stat = nf90_create('pot_solve_failstate_plane'//plane_c//'.nc', &
                                    NF90_NETCDF4+NF90_CLOBBER, nf90_id)
                                    
            call handle_error_netcdf(nf90_stat, __LINE__, __FILE__)
            call write_netcdf_helmholtz(nf90_id, mesh_cano, bnddescr_actual_pot_core, &
                                        bnddescr_actual_pot_wall, bnddescr_actual_pot_dome, &
                                        bnddescr_actual_pot_out, &
                                        co, lambda, xi, &
                                        rhs, &
                                        guess=potv, &
                                        sol=potv, &
                                        hcsr_write_on=.true.)
            nf90_stat = nf90_close(nf90_id)
            call handle_error_netcdf(nf90_stat, __LINE__, __FILE__)
        endif
                
        ! Extrapolate ghost points
        call perf_start('../../../../pot_solve_ghost')

        !$omp parallel default(none) shared(mesh_cano, potv)
        call extrapolate_ghost_points(mesh_cano, potv)
        !$omp end parallel

        call perf_stop('../../../../pot_solve_ghost')
            
    end subroutine


    subroutine parallel_polarisation_advection(comm_handler, equi, mesh_cano, mesh_stag, map, penalisation_cano, &
                  boundaries, co, pot, ti, ne, logne, uparv_cano, dvort)
        !! Computes the parallel advection part in the vorticity equation
        type(comm_handler_t), intent(in) :: comm_handler
        !! Communicators
        class(equilibrium_t), intent(in) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(in) :: mesh_cano
        !! Mesh (canonical)
        type(mesh_cart_t), intent(in) :: mesh_stag
        !! Mesh (staggered)
        type(parallel_map_t), intent(in) :: map
        !! Parallel map
        type(penalisation_t), intent(in) :: penalisation_cano
        !! Penalisation (canonical)
        type(boundaries_braginskii_t), intent(in) :: boundaries
        !! Boundary information for the BRAGINSKII model
        real(GP), dimension(:), intent(in) :: co
        !! Polarisation coefficient n/B^2, with n advanced to t+1
        type(variable_t), intent(in) :: pot
        !! Electrostatic potential
        type(variable_t), intent(in) :: ti
        !! Ion temperature
        type(variable_t), intent(in) :: ne
        !! Electron density
        type(variable_t), intent(in) :: logne
        !! Logarithm of the electron density
        real(GP), dimension(mesh_cano%get_n_points()), intent(in) :: uparv_cano
        !! Values of parallel velocity mapped to canonical grid
        real(GP), dimension(mesh_cano%get_n_points()), intent(inout) :: dvort
        !! Change of vorticity

        real(GP), dimension(mesh_cano%get_n_points()) :: dpar_pot, dpar_logne, dpar_pri
        real(GP), dimension(mesh_cano%get_n_points()) :: aux1, aux2, aux3, aux4
        real(GP), dimension(mesh_cano%get_n_points()) :: co_pot, co_dia, co_logne
        real(GP) :: phi, x, y, fac, dia_fac, pen_fac, xi
        integer :: ki, kb, kg, l

        type(variable_t) :: par_grad_pot_stag, par_grad_logne_stag, pri, par_grad_pri_stag

        phi = mesh_cano%get_phi()
        fac = - rhos**2 * swb_vort_paradv
        dia_fac = tratio * swb_vort_dia

        ! Compute parallel gradients of pot, logne and ion pressure on canonical mesh
        if (equi%is_axisymmetric()) then
            ! For axisymmetric equilibria use method that avoids MPI communication
            call pri%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)

            !$omp parallel default(none) private(l) &
            !$omp shared(mesh_cano, map, pot, logne, ne, ti, pri, &
            !$omp        dpar_pot, dpar_logne, dpar_pri, aux1, aux2, aux3, aux4)
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                ! aux1 = parallel_gradient(potv) on k+1/2, aux2 = parallel_gradient(potv) on k-1/2
                aux1(l) = map%par_grad_single(pot%vals, pot%hfwd, l)
                aux2(l) = map%par_grad_single(pot%hbwd, pot%vals, l)
                ! aux3 = parallel_gradient(lognev) on k+1/2, aux4 = parallel_gradient(lognev) on k-1/2
                aux3(l) = map%par_grad_single(logne%vals, logne%hfwd, l)
                aux4(l) = map%par_grad_single(logne%hbwd, logne%vals, l)
            enddo
            !$omp end do
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                ! on full grid !
                dpar_pot(l) = map%flat_stag_to_cano_single(aux1, aux2, l)
                dpar_logne(l) = map%flat_stag_to_cano_single(aux3, aux4, l)
            enddo
            !$omp end do
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                pri%vals(l) = ne%vals(l) * ti%vals(l)   ! pri = ion pressure
                aux1(l) = ne%hfwd(l) * ti%hfwd(l) ! pri halo forward
                aux2(l) = ne%hbwd(l) * ti%hbwd(l) ! pri halo backward
            enddo
            !$omp end do
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                ! aux3 = parallel_gradient(pri) on k+1/2, aux4 = parallel_gradient(pri) on k-1/2
                aux3(l) = map%par_grad_single(pri%vals, aux1, l)
                aux4(l) = map%par_grad_single(aux2, pri%vals, l)
            enddo
            !$omp end do
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                dpar_pri(l) = map%flat_stag_to_cano_single(aux3, aux4, l)
            enddo
            !$omp end do
            !$omp end parallel

        else

            call pri%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)
            call pri%create_halos(comm_handler, 1, 0)
            !$omp parallel default(none) private(l) shared(mesh_cano, ne, ti, pri)
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                pri%vals(l) = ne%vals(l) * ti%vals(l)
            enddo
            !$omp end do
            !$omp end parallel

            call pri%start_comm_halos(comm_handler)
            call pri%finalize_comm_halos(comm_handler)

            call par_grad_pot_stag%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)
            call par_grad_pot_stag%create_halos(comm_handler, 0, 1)

            call par_grad_logne_stag%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)
            call par_grad_logne_stag%create_halos(comm_handler, 0, 1)

            call par_grad_pri_stag%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)
            call par_grad_pri_stag%create_halos(comm_handler, 0, 1)

            !$omp parallel default(none) &
            !$omp shared(map, pot, par_grad_pot_stag, logne, par_grad_logne_stag, pri, par_grad_pri_stag)
            call map%par_grad(pot%vals, pot%hfwd, par_grad_pot_stag%vals)
            call map%par_grad(logne%vals, logne%hfwd, par_grad_logne_stag%vals)
            call map%par_grad(pri%vals, pri%hfwd, par_grad_pri_stag%vals)
            !$omp end parallel

            call par_grad_pot_stag%start_comm_halos(comm_handler)
            call par_grad_pot_stag%finalize_comm_halos(comm_handler)

            call par_grad_logne_stag%start_comm_halos(comm_handler)
            call par_grad_logne_stag%finalize_comm_halos(comm_handler)

            call par_grad_pri_stag%start_comm_halos(comm_handler)
            call par_grad_pri_stag%finalize_comm_halos(comm_handler)

            !$omp parallel default(none) &
            !$omp shared(map, par_grad_pot_stag, dpar_pot, par_grad_logne_stag, dpar_logne, &
            !$omp        par_grad_pri_stag, dpar_pri )
            call map%flat_stag_to_cano(par_grad_pot_stag%vals, par_grad_pot_stag%hbwd, dpar_pot)
            call map%flat_stag_to_cano(par_grad_logne_stag%vals, par_grad_logne_stag%hbwd, dpar_logne)
            call map%flat_stag_to_cano(par_grad_pri_stag%vals, par_grad_pri_stag%hbwd, dpar_pri)
            !$omp end parallel

        endif

        ! Assemble parallel advection term
        !$omp parallel default(none) &
        !$omp private(ki, kb, kg, l, x, y, pen_fac, xi) &
        !$omp shared(equi, mesh_cano, penalisation_cano, phi, fac, dia_fac, &
        !$omp        co, ne, uparv_cano, pri, co_pot, co_dia, co_logne, &
        !$omp        dpar_pot, dpar_logne, dpar_pri, aux1, aux2, aux3, dvort)
        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            x = mesh_cano%get_x(l)
            y = mesh_cano%get_y(l)
            co_pot(l) = co(l) * uparv_cano(l) ! polarisation coefficient with upar
            co_dia(l) = co_pot(l) / ne%vals(l) * dia_fac
            co_logne(l) = co_dia(l) * dpar_logne(l)
        enddo
        !$omp end do
        !$omp do
        do kb = 1, mesh_cano%get_n_points_boundary()
            l = mesh_cano%boundary_indices(kb)
            co_pot(l) = co(l) * uparv_cano(l) ! polarisation coefficient with upar
            co_dia(l) = co_pot(l) / ne%vals(l) * dia_fac
            co_logne(l) = co_dia(l) * dpar_logne(l)
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh_cano%get_n_points_ghost()
            l = mesh_cano%ghost_indices(kg)
            co_pot(l) = co(l) * uparv_cano(l)
            co_dia(l) = co_pot(l) / ne%vals(l) * dia_fac
            co_logne(l) = co_dia(l) * dpar_logne(l)
        enddo
        !$omp end do

        !$omp do
        do ki = 1, mesh_cano%get_n_points_inner()
            l = mesh_cano%inner_indices(ki)
            x = mesh_cano%get_x(l)
            y = mesh_cano%get_y(l)
            xi = fac / equi%jacobian(x, y, phi)
            aux1(l) = helmholtz_single_inner(mesh_cano, dpar_pot, l, co_pot, xi)
            aux2(l) = helmholtz_single_inner(mesh_cano, dpar_pri, l, co_dia, xi)
            aux3(l) = helmholtz_single_inner(mesh_cano, pri%vals, l, co_logne, xi)
            pen_fac = 1.0_GP - penalisation_cano%get_charfun_val(ki)
            dvort(l) = dvort(l) + pen_fac * ( - aux1(l) - aux2(l) + aux3(l) )
        enddo
        !$omp end do
        !$omp end parallel

    end subroutine

end module