neutrals_module_m.f90 Source File


Contents

Source Code


Source Code

module neutrals_module_m
    !! Implementation of NEUTRALS model
    use MPI
    use comm_handler_m, only : comm_handler_t
    use precision_grillix_m, only : GP, GP_NAN, GP_EPS
    use parallel_map_m, only : parallel_map_t
    use penalisation_m, only : penalisation_t
    use variable_m, only: variable_t
    use multistep_m, only : karniadakis_t, multistep_storage_t
    use inplane_operators_m, only : inplane_operators_t
    use boundaries_neutrals_m, only : boundaries_neutrals_t
    use penalisation_neutrals_m, only: neutrals_dens_penalisation, &
                                       neutrals_parmom_penalisation, &
                                       neutrals_pressure_penalisation
    use mms_neutrals_m, only : mms_sol_neutrals_dens, mms_source_neutrals_dens, &
                             mms_sol_neutrals_parmom, mms_source_neutrals_parmom, &
                             mms_sol_neutrals_pressure, mms_source_neutrals_pressure
    use rate_coeff_neutrals_m, only: rate_coeff_neutrals_t
    use diff_coeff_neutrals_m, only: eval_diff_coeff
    use coronal_impurities_m, only: coronal_impurity_t
    use equilibrium_storage_m, only : equilibrium_storage_t
    use error_handling_grillix_m, only: handle_error, error_info_t
    use status_codes_grillix_m, only : GRILLIX_ERR_OTHER
    use descriptors_braginskii_m, only : BND_TYPE_DIRICHLET_ZERO, BND_TYPE_NEUMANN
    use boundaries_perp_m, only : set_perpbnds, extrapolate_ghost_points
    use parallel_target_flux_m, only : parallel_target_flux_t
    use perp_bnd_flux_m, only : perp_bnd_flux_t, mirrorset_polyperpbnd_points
    use recycling_m, only : compute_recycling_source
    use screen_io_m, only :  get_stdout

    ! From PARALLAX
    use perf_m, only : perf_start, perf_stop
    use equilibrium_m, only : equilibrium_t
    use mesh_cart_m, only : mesh_cart_t
    use helmholtz_solver_m, only : helmholtz_solver_t
    use boundaries_perp_m, only : extrapolate_ghost_points
    use helmholtz_operator_m, only : helmholtz_single_inner
    use helmholtz_solver_m, only : helmholtz_solver_t
    use helmholtz_solver_factory_m, only : parameters_helmholtz_solver_factory, &
                                           helmholtz_solver_factory
    use sources_external_m, only : source_gaussian_t

    ! Parameters
    use params_feature_selection_m, only : &
        mms_on, mms_neutrals_temp_coeff
    use params_tstep_m, only : &
        dtau, tstep_order
    use params_brag_model_m, only : &
        rhos, tratio
    use params_brag_boundaries_parpen_m, only : &
        pen_epsinv
    use params_neut_model_m, only : &
        ne_ref, te_ref, k_ref, w_ref, s_cx, pot_iz, &
        pardiss_coeff_dens, pardiss_coeff_parmom, pardiss_coeff_pressure, &
        evolve_pressure, apply_neumann_mirror, src_rcy_temp, cx_cooling_thresh, cx_cooling_thresh_min, &
        diff_co_smoothing
    use params_impy_coronal_m, only : &
        impy_concentration, impy_path, impy_tab_length
    use params_neut_floors_m, only : &
        floor_dens, floor_temp, rho_min_for_parmom_advection
    use params_neut_data_paths_m, only : &
        path_k_iz, path_k_rec, path_w_iz, path_p_rec
    use params_neut_switches_m
    use params_neut_boundaries_perp_m, only : &
        bnddescr_neut_dens_core, bnddescr_neut_dens_wall, &
        bnddescr_neut_dens_dome, bnddescr_neut_dens_out, &
        bnddescr_neut_parmom_core, bnddescr_neut_parmom_wall, &
        bnddescr_neut_parmom_dome, bnddescr_neut_parmom_out, &
        bnddescr_neut_pressure_core, bnddescr_neut_pressure_wall, &
        bnddescr_neut_pressure_dome, bnddescr_neut_pressure_out

    implicit none

    ! Auxiliary variables
    type(rate_coeff_neutrals_t), private, save :: k_iz
    ! Ionization rate coefficient object
    type(rate_coeff_neutrals_t), private, save :: k_rec
    ! Recombination rate coefficient object
    type(rate_coeff_neutrals_t), private, save :: w_iz
    ! Ionization cooling rate coefficient object
    type(rate_coeff_neutrals_t), private, save :: p_rec
    ! Recombination cooling rate coefficient object,
    ! evaluated as w_rec = (p_rec - 13.6 eV / te_ref * k_rec)
    type(coronal_impurity_t), private, save :: l_imp
    ! Coronal impurity (nitrogen) radiation rate coefficient object

    type(variable_t), private, save :: neutrals_temp
    ! Neutrals temperature
    type(variable_t), public, save :: diff_co
    ! Coefficient for diffusion term D_N / TN = 1 / (ne k_cx)
    real(GP), dimension(:), allocatable, public, save :: diff_co_presmooth
    ! Coffecient prior to additional smoothing (to be handled by diagnostics)

    type(variable_t), private, save :: src_freq_ne
    ! Density source frequency, i.e. plasma density source divided by plasma density
    type(variable_t), private, save :: src_parmom_c1, src_parmom_c2
    ! Source coefficients for neutrals parallel momentum equation
    real(GP), dimension(:), allocatable, private, save :: src_neutrals_dens
    ! Neutrals density source
    real(GP), dimension(:), allocatable, private, save :: src_neutrals_parmom
    ! Parallel momentum source
    real(GP), dimension(:), allocatable, private, save :: src_neutrals_pressure
    ! Neutrals temperature source

    type(variable_t), private, save :: gradpar_neutrals_dens
    ! Parallel gradient of neutrals density
    type(variable_t), private, save :: divbpar_neutrals_parmom
    ! Parallel divergence of neutrals parallel momentum
    type(variable_t), private, save :: par_flux_parmom
    ! Flux of parallel momentum in the neutrals parallel momentum equation

    type(variable_t), private, save :: neutrals_vpar
    ! Neutrals parallel velocity on the staggered mesh
    type(variable_t), private, save :: gradpar_neutrals_pressure
    ! Parallel gradient of neutrals pressure
    type(variable_t), private, save :: par_flux_pressure
    ! Pressure flux in the neutrals pressure equation
    type(variable_t), private, save :: neutrals_temp_guess
    ! Extrapolated neutrals temperature at time t+1
    type(multistep_storage_t), private, save :: neutrals_temp_store
    ! Storage of neutrals temperature values at time-steps t,t-1,t-2,...

    real(GP), dimension(:), allocatable, public, save :: src_floor_neutrals_dens
    ! Indirect neutrals density source created by applying floor value
    ! Is written out as diagnostic, but must be measured during neutrals timestep

    real(GP), dimension(:,:), allocatable, private, save :: ki_nsegs
    !! Inner mesh index values associated with each point on boundary polygon
    integer, private, save :: cnt = 0
    !! Internal counter of timestep calls, for tracking when time extrapolation is valid

    type, public :: neutrals_module_t
        !! Neutrals module responsible for evaluating neutrals-plasma source
        !! and time evolution of neutrals quantities
        type(source_gaussian_t) :: gas_puff
    contains
        procedure, public :: init => init_neutrals
        procedure, public :: eval_sources => eval_sources_neutrals
        procedure, public :: timestep => timestep_neutrals
        procedure, private :: setup_implicit_solve
        procedure, private :: neumann_mirror
    end type

contains

    subroutine init_neutrals(self, comm_handler, equi, mesh_cano, mesh_stag, map, penalisation_cano, &
                             ne, neutrals_dens, neutrals_parmom, neutrals_pressure, isnaps_neutrals, tau)
        !! Setup data structures for neutrals timestepping
        class(neutrals_module_t), intent(inout) :: self
        !! Instance of class
        type(comm_handler_t), intent(in) :: comm_handler
        !! Communicators
        class(equilibrium_t), intent(inout) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(inout) :: mesh_cano
        !! Mesh (canonical) within poloidal plane
        type(mesh_cart_t), intent(inout) :: 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(variable_t), intent(inout) :: ne
        !! Electron density
        type(variable_t), intent(inout) :: neutrals_dens
        !! Neutrals density
        type(variable_t), intent(inout) :: neutrals_parmom
        !! Neutrals parallel momentum
        type(variable_t), intent(inout) :: neutrals_pressure
        !! Neutrals pressure
        integer, intent(in) :: isnaps_neutrals
        !! Snapshot number for neutrals
        real(GP), intent(in) :: tau
        !! Time

        integer :: l, kt

        ! Load rate coefficients
        call k_iz%create(path_k_iz, ne_ref, &
                            te_ref, k_ref)
        call k_rec%create(path_k_rec, ne_ref, &
                            te_ref, k_ref)
        ! Cooling rate coefficients use a different rate coefficient reference value
        call w_iz%create(path_w_iz, ne_ref,    &
                            te_ref, w_ref)
        call p_rec%create(path_p_rec, ne_ref,  &
                            te_ref, w_ref)
        ! Load coronal impurity radiation rate coefficient (if impy_concentration given)
        if (impy_concentration > GP_EPS) then
            call l_imp%create(impy_path, impy_tab_length, &
                        te_ref, w_ref, impy_concentration)
        end if

        ! Allocate work variables
        call neutrals_temp%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)
        call neutrals_temp_guess%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)
        call diff_co%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)

        call gradpar_neutrals_dens%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)
        call divbpar_neutrals_parmom%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)
        call par_flux_parmom%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)

        call neutrals_vpar%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)
        call gradpar_neutrals_pressure%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)
        call par_flux_pressure%init(comm_handler, mesh_stag%get_n_points(), staggered=.true.)

        call src_freq_ne%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)
        call src_parmom_c1%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)
        call src_parmom_c2%init(comm_handler, mesh_cano%get_n_points(), staggered=.false.)

        allocate(src_neutrals_dens(mesh_cano%get_n_points_inner()), source=0.0_GP)
        allocate(src_neutrals_parmom(mesh_stag%get_n_points_inner()), source=0.0_GP)
        allocate(src_neutrals_pressure(mesh_cano%get_n_points_inner()), source=0.0_GP)
        allocate(src_floor_neutrals_dens(mesh_cano%get_n_points()), source=0.0_GP)

        allocate(diff_co_presmooth(mesh_cano%get_n_points()), source=0.0_GP)

        ! Setup halo structures
        call neutrals_temp%create_halos(comm_handler, 1, 1)
        call neutrals_temp_guess%create_halos(comm_handler, 1, 1)
        call diff_co%create_halos(comm_handler, 1, 1)

        call gradpar_neutrals_dens%create_halos(comm_handler, 0, 1)
        call divbpar_neutrals_parmom%create_halos(comm_handler, 1, 0)
        call par_flux_parmom%create_halos(comm_handler, 1, 0)

        call neutrals_vpar%create_halos(comm_handler, 0, 1)
        call gradpar_neutrals_pressure%create_halos(comm_handler, 0, 1)
        call par_flux_pressure%create_halos(comm_handler, 0, 1)

        call src_freq_ne%create_halos(comm_handler, 1, 0)
        call src_parmom_c1%create_halos(comm_handler, 1, 0)
        call src_parmom_c2%create_halos(comm_handler, 1, 0)

        ! In case the initial state (e.g. from snapshot) contains zeros, apply floor
        !$omp parallel default(none) private(l) &
        !$omp          shared(mesh_cano, neutrals_dens, floor_dens, neutrals_pressure, floor_temp)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            neutrals_dens%vals(l) = max(neutrals_dens%vals(l), floor_dens)
            neutrals_pressure%vals(l) = max(neutrals_pressure%vals(l), floor_dens*floor_temp)
        enddo
        !$omp end do
        !$omp end parallel

        ! Initialize multistep storage
        call neutrals_temp_store%init_storage(mesh_cano%get_n_points(), tstep_order-1, 1)
        !$omp parallel default(none) private(l, kt) &
        !$omp          shared(mesh_cano, tstep_order, &
        !$omp                 neutrals_dens, neutrals_pressure, &
        !$omp                 neutrals_temp_store)
        do kt = 1, tstep_order-1
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                neutrals_temp_store%vstore(l, kt, 1) = neutrals_pressure%vals(l) / neutrals_dens%vals(l)
            enddo
            !$omp end do
        enddo
        !$omp end parallel

        call self%gas_puff%init('params_neutrals.in')
        if ( self%gas_puff%source_gaussian_on ) then
            call self%gas_puff%set_coeffs(comm_handler, mesh_cano, map, penalisation_cano)
            call self%gas_puff%display()
        endif

    end subroutine

    subroutine eval_sources_neutrals(self, comm_handler, equi, equi_on_cano, mesh_cano, mesh_stag, map, &
                                ne, pot, upar, te, ti, neutrals_dens, neutrals_parmom, neutrals_pressure, &
                                src_ne, src_upar, src_te, src_ti, src_vort)
        !! Evaluate source terms to pass to Braginskii model
        class(neutrals_module_t), intent(inout) :: self
        !! Instance of class
        type(comm_handler_t), intent(in) :: comm_handler
        !! Communicators
        class(equilibrium_t), intent(inout) :: equi
        !! Equilibrium
        class(equilibrium_storage_t), intent(inout) :: equi_on_cano
        !! Equilibrim quantities on canonical mesh
        type(mesh_cart_t), intent(inout) :: mesh_cano
        !! Mesh (canonical) within poloidal plane
        type(mesh_cart_t), intent(inout) :: mesh_stag
        !! Mesh (staggered) within poloidal plane
        type(parallel_map_t), intent(in) :: map
        !! Parallel map
        type(variable_t), intent(inout) :: ne
        !! Plasma density at timestep t
        type(variable_t), intent(in) :: pot
        !! Electrostatic potential
        type(variable_t), intent(inout) :: upar
        !! Parallel ion velocity
        type(variable_t), intent(in) :: te
        !! Electron temperature at timestep t
        type(variable_t), intent(in) :: ti
        !! Ion temperature at timestep t
        type(variable_t), intent(inout) :: neutrals_dens
        !! Neutrals density, on input at t, on output advanced to t+1
        type(variable_t), intent(inout) :: neutrals_parmom
        !! Neutrals parallel momentum, on input at t, on output advanced to t+1
        type(variable_t), intent(inout) :: neutrals_pressure
        !! Neutrals pressure, on input at t, on output advanced to t+1
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(inout) :: src_ne
        !! Plasma particle source, to which neutrals model contribution is added
        real(GP), dimension(mesh_stag%get_n_points_inner()), intent(inout) :: src_upar
        !! Parallel momentum source, to which neutrtype(comm_handler_t), intent(in) :: comm_handler
        !! Communicatorsals model contribution is added
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(inout) :: src_te
        !! Electron temperature source, to which neutrals model contribution is added
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(inout) :: src_ti
        !! Ion temperature source, to which neutrals model contribution is added
        real(GP), dimension(mesh_cano%get_n_points_inner()), intent(inout) :: src_vort
        !! Vorticity source, to which neutrals model contribution is added

        integer :: l, ki
        real(GP) :: x, y, phi_cano, k_iz_l, k_rec_l, k_cx_l, cool_iz_l, cool_rec_l, cool_impurity_l, &
                    src_freq_ne_stag, c1_stag, c2_stag, ne_stag, upar_cano_l, neutrals_vpar_cano_l, &
                    upar_vpar_friction, cx_equi_ti, cx_equi_pn
        real(GP), dimension(mesh_cano%get_n_points()) :: pi, co_vortpi, co_vortpot
        ! Ion pressure (ne * Ti), coefficients related to vorticity source

        phi_cano = mesh_cano%get_phi()

        !$omp parallel default(none) &
        !$omp          private(l, x, y, k_iz_l, k_rec_l, k_cx_l) &
        !$omp          shared(equi, mesh_cano, phi_cano, equi_on_cano, &
        !$omp                 k_iz, k_rec, tratio, s_cx, &
        !$omp                 swn_iz, swn_rec,  &
        !$omp                 ne, te, ti, pi, neutrals_dens, &
        !$omp                 co_vortpi, co_vortpot, src_parmom_c1, src_parmom_c2, src_freq_ne)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            x = mesh_cano%get_x(l)
            y = mesh_cano%get_y(l)
            k_iz_l = swn_iz * k_iz%eval(te%vals(l), ne%vals(l))
            k_rec_l = swn_rec * k_rec%eval(te%vals(l), ne%vals(l))
            k_cx_l = 2.93_GP * s_cx * sqrt(ti%vals(l) * tratio )
            ! Frequency of density change (needed on staggered grid)
            src_freq_ne%vals(l) = k_iz_l * neutrals_dens%vals(l) - k_rec_l * ne%vals(l)
            ! Ion pressure part of vorticity source
            pi(l) = ne%vals(l) * ti%vals(l) * tratio
            ! Vorticity source coefficients
            co_vortpi(l) = (k_iz_l + k_cx_l) * neutrals_dens%vals(l) * equi%jacobian(x, y, phi_cano) / equi_on_cano%absb(l)**2
            co_vortpot(l) = co_vortpi(l) * ne%vals(l)
            ! Parallel momentum source coefficients
            src_parmom_c1%vals(l) = (k_rec_l * ne%vals(l) + k_cx_l * neutrals_dens%vals(l))
            src_parmom_c2%vals(l) = (k_iz_l + k_cx_l)
        enddo
        !$omp end do
        !$omp end parallel

        call src_freq_ne%start_comm_halos(comm_handler)
        call src_freq_ne%finalize_comm_halos(comm_handler)

        call neutrals_parmom%start_comm_halos(comm_handler)
        call neutrals_parmom%finalize_comm_halos(comm_handler)

        call src_parmom_c1%start_comm_halos(comm_handler)
        call src_parmom_c1%finalize_comm_halos(comm_handler)

        call src_parmom_c2%start_comm_halos(comm_handler)
        call src_parmom_c2%finalize_comm_halos(comm_handler)

        call ne%start_comm_halos(comm_handler)
        call ne%finalize_comm_halos(comm_handler)

        call upar%start_comm_halos(comm_handler)
        call upar%finalize_comm_halos(comm_handler)

        !$omp parallel default(none) &
        !$omp          private(l, ki, x, y, &
        !$omp                  k_iz_l, k_rec_l, k_cx_l, cool_iz_l, cool_rec_l, cool_impurity_l, &
        !$omp                  src_freq_ne_stag, c1_stag, c2_stag, ne_stag, &
        !$omp                  upar_cano_l, neutrals_vpar_cano_l, upar_vpar_friction, cx_equi_ti, cx_equi_pn) &
        !$omp          shared(equi, mesh_cano, phi_cano, equi_on_cano, mesh_stag, map, &
        !$omp                 rhos, tratio, k_iz, k_rec, w_iz, p_rec, l_imp, rho_min_for_parmom_advection, &
        !$omp                 s_cx, pot_iz, te_ref, impy_concentration, &
        !$omp                 ne, pot, upar, te, ti, pi, &
        !$omp                 neutrals_dens, neutrals_parmom, neutrals_temp, neutrals_pressure, &
        !$omp                 co_vortpi, co_vortpot, src_parmom_c1, src_parmom_c2, &
        !$omp                 src_ne, src_te, src_ti, src_vort, src_upar, src_freq_ne, &
        !$omp                 src_neutrals_dens, src_neutrals_parmom, src_neutrals_pressure, &
        !$omp                 swn_iz, swn_rec, swn_src_parmom, swn_src_pressure, &
        !$omp                 swn_src_te, swn_src_ti, swn_src_vort, swn_src_upar, swn_cx_equi_ti, swn_cx_equi_pn, &
        !$omp                 cx_cooling_thresh, cx_cooling_thresh_min)
        !$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)
            ! Rate coefficients
            k_iz_l = swn_iz * k_iz%eval(te%vals(l), ne%vals(l))
            k_rec_l = swn_rec * k_rec%eval(te%vals(l), ne%vals(l))
            k_cx_l = 2.93_GP * s_cx * sqrt(ti%vals(l) * tratio )
            cool_iz_l = swn_iz * w_iz%eval(te%vals(l), ne%vals(l))
            cool_rec_l = swn_rec * (p_rec%eval(te%vals(l), ne%vals(l)) - pot_iz / te_ref * k_rec_l)
            if (impy_concentration > GP_EPS) then
                cool_impurity_l = l_imp%eval(te%vals, l)
            end if
            ! Neutrals temperature
            neutrals_temp%vals(l) = neutrals_pressure%vals(l) / neutrals_dens%vals(l)
            ! Mappings
            upar_cano_l = map%flat_stag_to_cano_single(upar%vals, upar%hbwd, l)
            neutrals_vpar_cano_l = map%flat_stag_to_cano_single(neutrals_parmom%vals, neutrals_parmom%hbwd, l) / neutrals_dens%vals(l)

            upar_vpar_friction = ( upar_cano_l - neutrals_vpar_cano_l )**2

            if (equi_on_cano%rho(l) < rho_min_for_parmom_advection) then
                ! Towards the plasma core, N is close to zero, and computing
                ! parallel neutrals velocity related terms (where a division
                ! by N occurs) can lead to numerical problems.
                upar_vpar_friction = 0.0_GP
            endif

            if ( (ti%vals(l) < cx_cooling_thresh) .and. (ti%vals(l) > cx_cooling_thresh_min) ) then
                cx_equi_ti = k_cx_l * ( 0.0_GP - neutrals_dens%vals(l) * ti%vals(l) )
                cx_equi_pn = - k_cx_l * ( neutrals_pressure%vals(l) - neutrals_dens%vals(l) * ti%vals(l) )
            else
                cx_equi_ti = k_cx_l * ( neutrals_pressure%vals(l) - neutrals_dens%vals(l) * ti%vals(l) )
                cx_equi_pn = - cx_equi_ti
            endif

            ! Compute source terms
            src_neutrals_dens(ki) = - src_freq_ne%vals(l) * ne%vals(l)

            src_ne(ki) = src_ne(ki) - src_neutrals_dens(ki)

            src_te(ki) = src_te(ki) - swn_src_te * ( &
                            + te%vals(l) * src_freq_ne%vals(l) &
                            - 2.0_GP/3.0_GP * (cool_iz_l * neutrals_dens%vals(l) + cool_rec_l * ne%vals(l)))
            if (impy_concentration > 0.0_GP) then
                src_te(ki) = src_te(ki) - swn_src_te * ( &
                                    2.0_GP/3.0_GP * (cool_impurity_l * ne%vals(l) * impy_concentration))
            end if

            src_ti(ki) = src_ti(ki) + swn_src_ti * ( &
                                + k_iz_l * neutrals_pressure%vals(l) &
                                - k_iz_l * neutrals_dens%vals(l) * ti%vals(l) &
                                + swn_cx_equi_ti * cx_equi_ti &
                                + neutrals_dens%vals(l) * (k_iz_l + k_cx_l) / 3.0_GP * upar_vpar_friction )

            src_vort(ki) = src_vort(ki) + swn_src_vort * ( &
                              helmholtz_single_inner(mesh_cano, pot%vals, l, co=co_vortpot, xiv=1.0_GP, lambdav=0.0_GP) &
                            + helmholtz_single_inner(mesh_cano, pi, l, co=co_vortpi, xiv=1.0_GP, lambdav=0.0_GP)  &
                           ) * (rhos)**2 / equi%jacobian(x, y, phi_cano)

            src_neutrals_pressure(ki) = swn_src_pressure * ( &
                                        + k_rec_l * ne%vals(l)**2 * ti%vals(l)  &
                                        - k_iz_l * ne%vals(l) * neutrals_pressure%vals(l) &
                                        + swn_cx_equi_pn * ne%vals(l) * cx_equi_pn &
                                        + ne%vals(l) / 3.0_GP * ( ne%vals(l) * k_rec_l + neutrals_dens%vals(l) * k_cx_l ) &
                                           * upar_vpar_friction &
                                        )
        enddo
        !$omp end do
        !$omp do
        do ki = 1, mesh_stag%get_n_points_inner()
            l = mesh_stag%inner_indices(ki)
            ! Mappings
            src_freq_ne_stag = map%lift_cano_to_stag_single(src_freq_ne%vals, src_freq_ne%hfwd, l)
            c1_stag = map%lift_cano_to_stag_single(src_parmom_c1%vals, src_parmom_c1%hfwd, l)
            c2_stag = map%lift_cano_to_stag_single(src_parmom_c2%vals, src_parmom_c2%hfwd, l)
            ne_stag = map%lift_cano_to_stag_single(ne%vals, ne%hfwd, l)

            src_neutrals_parmom(ki) = swn_src_parmom * ne_stag * (c1_stag * upar%vals(l) - c2_stag * neutrals_parmom%vals(l))

            src_upar(ki) = src_upar(ki) - swn_src_upar * upar%vals(l) * src_freq_ne_stag - src_neutrals_parmom(ki) / ne_stag
        enddo
        !$omp end do
        !$omp end parallel

    end subroutine

    subroutine timestep_neutrals(self, comm_handler, &
                                 equi, equi_on_cano, equi_on_stag, &
                                 mesh_cano, mesh_stag, &
                                 hsolver_cano, hsolver_stag, &
                                 map, &
                                 penalisation_cano, penalisation_stag, &
                                 parflux_cano, parflux_stag, &
                                 perp_bnd_flux_cano, perp_bnd_flux_stag, &
                                 opsinplane_cano, opsinplane_stag, &
                                 boundaries_neutrals, tau, &
                                 ne, pot, upar, jpar, apar_fluct, te, ti, &
                                 neutrals_dens, neutrals_parmom, neutrals_pressure, &
                                 tstep_neutrals_dens, tstep_neutrals_parmom, &
                                 tstep_neutrals_pressure, &
                                 sinfo, res, success_neutrals)
        !! Advances neutrals quantities from t to t+1
        class(neutrals_module_t), intent(inout) :: self
        !! Instance of class
        type(comm_handler_t), intent(in) :: comm_handler
        !! Communicators
        class(equilibrium_t), intent(inout) :: equi
        !! Equilibrium
        class(equilibrium_storage_t), intent(inout) :: equi_on_cano
        !! Equilibrim quantities on canonical mesh
        class(equilibrium_storage_t), intent(inout) :: equi_on_stag
        !! Equilibrim quantities on staggered mesh
        type(mesh_cart_t), intent(inout) :: mesh_cano
        !! Mesh (canonical) within poloidal plane
        type(mesh_cart_t), intent(inout) :: mesh_stag
        !! Mesh (staggered) within poloidal plane
        class(helmholtz_solver_t), intent(inout) :: hsolver_cano
        !! Elliptic (2D) solver on canonical mesh
        class(helmholtz_solver_t), intent(inout) :: hsolver_stag
        !! Elliptic (2D) solver on staggered mesh
        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(parallel_target_flux_t), intent(in) :: parflux_cano
        !! Parallel target flux utility (canonical)
        type(parallel_target_flux_t), intent(in) :: parflux_stag
        !! Parallel target flux utility (staggered)
        type(perp_bnd_flux_t), intent(in) :: perp_bnd_flux_cano
        !! Perpendicular boundary flux utility (canonical)
        type(perp_bnd_flux_t), intent(in) :: perp_bnd_flux_stag
        !! Perpendicular boundary flux utility (staggered)
        type(inplane_operators_t), intent(inout) :: opsinplane_cano
        !! In-plane operators (canonical)
        type(inplane_operators_t), intent(inout) :: opsinplane_stag
        !! In-plane operators (staggered)
        type(boundaries_neutrals_t), intent(inout) :: boundaries_neutrals
        !! Boundary information for the NEUTRALS model
        real(GP), intent(in) :: tau
        !! Time
        type(variable_t), intent(in) :: ne
        !! Plasma density at timestep t
        type(variable_t), intent(in) :: pot
        !! Electrostatic potential
        type(variable_t), intent(inout) :: upar
        !! Parallel ion velocity
        type(variable_t), intent(inout) :: jpar
        !! Parallel current
        type(variable_t), intent(inout) :: apar_fluct
        !! Fluctuation of apar used for flutter operators
        type(variable_t), intent(in) :: te
        !! Electron temperature at timestep t
        type(variable_t), intent(in) :: ti
        !! Ion temperature at timestep t
        type(variable_t), intent(inout) :: neutrals_dens
        !! Neutrals density, on input at t, on output advanced to t+1
        type(variable_t), intent(inout) :: neutrals_parmom
        !! Neutrals parallel momentum, on input at t, on output advanced to t+1
        type(variable_t), intent(inout) :: neutrals_pressure
        !! Neutrals pressure, on input at t, on output advanced to t+1
        type(karniadakis_t), intent(inout) :: tstep_neutrals_dens
        !! Time-step integrator for neutrals density
        type(karniadakis_t), intent(inout) :: tstep_neutrals_parmom
        !! Time-step integrator for neutrals parallel momentum
        type(karniadakis_t), intent(inout) :: tstep_neutrals_pressure
        !! Time-step integrator for neutrals pressure
        integer, intent(out), dimension(4) :: sinfo
        !! Info from elliptic solver
        real(GP), intent(out), dimension(4) :: res
        !! Residual of solution of elliptic solver
        logical, intent(out) :: success_neutrals
        !! Success flag for timestep

        integer :: ki
        ! Loop index running over inner grid points
        integer :: kb
        ! Loop index running over boundary points
        integer :: kg
        ! Loop index running over ghost points
        integer :: l
        ! Mesh index

        real(GP) :: tau_adv
        ! Time at timestep t+1
        real(GP) :: x, y, phi_cano, phi_stag
        ! Coordinates

        real(GP) :: pen_fac
        ! Penalisation factor

        real(GP), dimension(mesh_cano%get_n_points()) :: neutrals_dens_adv, dneutrals_dens
        ! Neutrals density explicit changerate and advanced in time at t+1
        real(GP), dimension(mesh_stag%get_n_points()) :: neutrals_parmom_adv, dneutrals_parmom
        ! Neutrals parallel momentum explicit changerate and advanced in time at t+1
        real(GP), dimension(mesh_cano%get_n_points()) :: neutrals_pressure_adv, dneutrals_pressure
        ! Neutrals pressure explicit changerate and advanced in time at t+1
        real(GP), dimension(mesh_cano%get_n_points_inner()) :: neutrals_dens_pen
        ! Penalisation values for neutrals density
        real(GP), dimension(mesh_stag%get_n_points_inner()) :: neutrals_parmom_pen
        ! Penalisation values for neutrals parallel momentum
        real(GP), dimension(mesh_cano%get_n_points_inner()) :: neutrals_pressure_pen
        ! Penalisation values for neutrals pressure

        real(GP), dimension(mesh_cano%get_n_points()) :: co_cano, rhs_cano
        real(GP), dimension(mesh_stag%get_n_points()) :: co_stag, rhs_stag

        real(GP), dimension(mesh_cano%get_n_points_inner()) :: lambda_cano, xi_cano
        real(GP), dimension(mesh_stag%get_n_points_inner()) :: lambda_stag, xi_stag
        ! Coefficients for implicit solve

        real(GP), dimension(mesh_stag%get_n_points()) :: neutrals_temp_guess_stag
        ! Neutrals temperature on staggered grid extrapolated at time t+1

        real(GP) :: diff_co_fwdavg = 0.0_GP
        ! Forward average of coefficient in toroidal diffusion, 0.5*(c_fwd + c)
        real(GP) :: diff_co_bwdavg = 0.0_GP
        ! Backward average of coefficient in toroidal diffusion, 0.5*(c + c_bwd)
        real(GP) :: diff_tor
        ! Toroidal diffusion term

        real(GP) :: hf2
        ! Mesh fine spacing squared

        real(GP) :: diss_par_neutrals_dens, diss_par_neutrals_parmom, diss_par_neutrals_pressure
        real(GP) :: par_grad_par_flux, par_grad_pressure, div_par_flux_pressure, pressure_div_vpar
        real(GP) :: grad_vpar_sqrd, pressure_vischeat, &
                    ddx_neutrals_vpar_cano, ddy_neutrals_vpar_cano
        ! Parallel operator auxiliaries

        real(GP), dimension(mesh_cano%get_n_points()) :: neutrals_parmom_cano
        ! Parallel neutrals momentum on canonical mesh, auxiliary field
        real(GP), dimension(mesh_cano%get_n_points()) :: neutrals_vpar_cano
        ! Parallel neutrals velocity on canonical mesh, auxiliary field
        real(GP), dimension(mesh_cano%get_n_points()) :: nvpar_cano
        ! Parallel electron flux on canonical mesh, auxiliary field
        real(GP), dimension(mesh_cano%get_n_points()) :: apar_fluct_cano
        ! Fluctuation of apar on canonical mesh, auxiliary field
        real(GP), dimension(mesh_cano%get_n_points()) :: src_rcy
        ! Recycling source rate
        real(GP), dimension(mesh_cano%get_n_points()) :: src_gas_puff
        ! Gas puffing/pumping source rate
        real(GP), dimension(mesh_cano%get_n_points()) :: diff_co_dens
        ! Effective diffusion coefficient in neutrals density equation
        real(GP), dimension(mesh_stag%get_n_points()) :: diff_co_parmom
        ! Effective diffusion coefficient in neutrals parallel momentum equation
        real(GP), dimension(mesh_cano%get_n_points()) :: diff_co_pressure
        ! Effective diffusion coefficient in neutrals pressure equation

        integer :: global_sinfo, ierr
        integer, dimension(mesh_cano%get_n_points_boundary()) :: bnd_descrs_nmn_cano
        integer, dimension(mesh_stag%get_n_points_boundary()) :: bnd_descrs_nmn_stag

        call perf_start('timestep_neutrals')

        phi_cano = mesh_cano%get_phi()
        phi_stag = mesh_stag%get_phi()
        tau_adv  = tau + dtau
        cnt = cnt + 1

        !$omp parallel default(none) private(l, kb) &
        !$omp          shared(mesh_cano, mesh_stag, bnd_descrs_nmn_cano, bnd_descrs_nmn_stag)
        !$omp do
        do kb = 1, mesh_cano%get_n_points_boundary()
            bnd_descrs_nmn_cano(kb) = BND_TYPE_NEUMANN
        enddo
        !$omp end do
        !$omp do
        do kb = 1, mesh_stag%get_n_points_boundary()
            bnd_descrs_nmn_stag(kb) = BND_TYPE_NEUMANN
        enddo
        !$omp end do
        !$omp end parallel

        ! Determine neutrals temperature ---------------------------------------
        call perf_start('../neutrals_temp')

        if (evolve_pressure) then
            !$omp parallel default(none) private(l) &
            !$omp          shared(mesh_cano, neutrals_dens, neutrals_pressure, floor_temp, &
            !$omp                 neutrals_temp, neutrals_temp_guess, neutrals_temp_store)
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                neutrals_temp%vals(l) = neutrals_pressure%vals(l) / neutrals_dens%vals(l)
                neutrals_temp_guess%vals(l) =   3.0_GP*neutrals_temp%vals(l) &
                                              - 3.0_GP*neutrals_temp_store%vstore(l,1,1) &
                                              + neutrals_temp_store%vstore(l,2,1)
                neutrals_temp_guess%vals(l) = max(neutrals_temp_guess%vals(l), floor_temp)
            enddo
            !$omp end do
            !$omp end parallel
        else
            !$omp parallel default(none) private(l, kb) &
            !$omp          shared(mesh_cano, ti, neutrals_temp, neutrals_temp_guess, floor_temp)
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                ! Apply temperature floor, if used
                neutrals_temp%vals(l) = max(ti%vals(l), floor_temp)
                neutrals_temp_guess%vals(l) = neutrals_temp%vals(l)
            enddo
            !$omp end do
            !$omp end parallel
        endif

        call perf_stop('../neutrals_temp')

        ! Communicate variables ------------------------------------------------
        call perf_start('../comm_neutrals_1')

        call neutrals_dens%start_comm_halos(comm_handler)
        call neutrals_dens%finalize_comm_halos(comm_handler)

        call neutrals_parmom%start_comm_halos(comm_handler)
        call neutrals_parmom%finalize_comm_halos(comm_handler)

        call neutrals_pressure%start_comm_halos(comm_handler)
        call neutrals_pressure%finalize_comm_halos(comm_handler)

        call neutrals_temp%start_comm_halos(comm_handler)
        call neutrals_temp%finalize_comm_halos(comm_handler)

        call neutrals_temp_guess%start_comm_halos(comm_handler)
        call neutrals_temp_guess%finalize_comm_halos(comm_handler)

        call upar%start_comm_halos(comm_handler)
        call upar%finalize_comm_halos(comm_handler)

        call jpar%start_comm_halos(comm_handler)
        call jpar%finalize_comm_halos(comm_handler)

        call apar_fluct%start_comm_halos(comm_handler)
        call apar_fluct%finalize_comm_halos(comm_handler)

        call perf_stop('../comm_neutrals_1')

        !$omp parallel do default(none) private(l) shared(mesh_stag, map, neutrals_temp_guess, neutrals_temp_guess_stag, floor_temp)
        do l = 1, mesh_stag%get_n_points()
            neutrals_temp_guess_stag(l) = map%lift_cano_to_stag_single(neutrals_temp_guess%vals, neutrals_temp_guess%hfwd, l)
            neutrals_temp_guess_stag(l) = max(neutrals_temp_guess_stag(l), floor_temp)
        enddo
        !$omp end parallel do

        if (mms_neutrals_temp_coeff) then
            ! Debugging feature: overwrite temperature guess with known mms solution in next timestep
            !$omp parallel default(none) private(l, x, y) &
            !$omp          shared(equi, mesh_cano, mesh_stag, &
            !$omp                 phi_cano, phi_stag, tau_adv, &
            !$omp                 neutrals_temp_guess, neutrals_temp_guess_stag)
            !$omp do
            do l = 1, mesh_cano%get_n_points()
                x = mesh_cano%get_x(l)
                y = mesh_cano%get_y(l)
                neutrals_temp_guess%vals(l) = mms_sol_neutrals_pressure(equi, x, y, phi_cano, tau_adv) / mms_sol_neutrals_dens(equi, x, y, phi_cano, tau_adv)
            enddo
            !$omp end do
            !$omp do
            do l = 1, mesh_stag%get_n_points()
                x = mesh_stag%get_x(l)
                y = mesh_stag%get_y(l)
                neutrals_temp_guess_stag(l) = mms_sol_neutrals_pressure(equi, x, y, phi_stag, tau_adv) / mms_sol_neutrals_dens(equi, x, y, phi_stag, tau_adv)
            enddo
            !$omp end do
            !$omp end parallel
        endif

        ! Compute auxilliary variables -----------------------------------------
        call perf_start('../aux_neutrals')

        ! Diffusion coefficient
        !   For toroidal diffusion (not used together with parallel flow!),
        !   neutrals_pressure must be communicated toroidally before entering the flux limiter!
        call eval_diff_coeff(equi, mesh_cano, map, ne, ti, neutrals_temp_guess, neutrals_pressure, diff_co%vals)

        ! Optionally, smoothen diffusion coefficient to increase numeric stability
        if (diff_co_smoothing > GP_EPS) then
            hf2 = mesh_cano%get_spacing_f()**2 
            !$omp parallel default(none) private(ki, kb, kg, l, x, y) &
            !$omp          shared(equi, mesh_cano, phi_cano, diff_co, diff_co_presmooth, diff_co_smoothing, hf2, &
            !$omp                 lambda_cano, xi_cano, co_cano, rhs_cano)
            !$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)
                diff_co_presmooth(l) = diff_co%vals(l)
                lambda_cano(ki) = 1.0_GP
                xi_cano(ki) = 1.0_GP / equi%jacobian(x, y, phi_cano)
                co_cano(l) = equi%jacobian(x, y, phi_cano) * diff_co_smoothing * hf2
                rhs_cano(l) = diff_co%vals(l)
            enddo
            !$omp end do
            !$omp do
            do kb = 1, mesh_cano%get_n_points_boundary()
                l = mesh_cano%boundary_indices(kb)
                x = mesh_cano%get_x(l)
                y = mesh_cano%get_y(l)
                diff_co_presmooth(l) = diff_co%vals(l)
                co_cano(l) = equi%jacobian(x, y, phi_cano) * diff_co_smoothing * hf2
                rhs_cano(l) = diff_co%vals(l)
            enddo
            !$omp end do
            !$omp do
            do kg = 1, mesh_cano%get_n_points_ghost()
                l = mesh_cano%ghost_indices(kg)
                x = mesh_cano%get_x(l)
                y = mesh_cano%get_y(l)
                diff_co_presmooth(l) = diff_co%vals(l)
                co_cano(l) = equi%jacobian(x, y, phi_cano) * diff_co_smoothing * hf2
                rhs_cano(l) = diff_co%vals(l)
            enddo
            !$omp end do
            !$omp end parallel

            call hsolver_cano%update(co_cano, lambda_cano, xi_cano, &
                                     BND_TYPE_DIRICHLET_ZERO, &
                                     BND_TYPE_DIRICHLET_ZERO, &
                                     BND_TYPE_DIRICHLET_ZERO, &
                                     BND_TYPE_DIRICHLET_ZERO)

            call hsolver_cano%solve(rhs_cano, diff_co%vals, res(4), sinfo(4)) 
        endif

        ! ---------------------------------------------------------------------------------------------------------

        !$omp parallel default(none) &
        !$omp          private(l) &
        !$omp          shared(equi, mesh_cano, mesh_stag, map, equi_on_cano, &
        !$omp                 k_iz, k_rec, ne, te, ti, upar, jpar, apar_fluct, nvpar_cano, apar_fluct_cano, &
        !$omp                 neutrals_dens, neutrals_parmom, neutrals_pressure, &
        !$omp                 neutrals_temp, neutrals_temp_guess, diff_co, &
        !$omp                 gradpar_neutrals_dens, divbpar_neutrals_parmom, par_flux_parmom, &
        !$omp                 neutrals_parmom_cano, neutrals_vpar, neutrals_vpar_cano, &
        !$omp                 gradpar_neutrals_pressure, par_flux_pressure)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            ! Parallel dynamics
            divbpar_neutrals_parmom%vals(l) = map%par_divb_single(neutrals_parmom%vals, neutrals_parmom%hbwd, l)
            neutrals_parmom_cano(l) = map%flat_stag_to_cano_single(neutrals_parmom%vals, neutrals_parmom%hbwd, l)
            par_flux_parmom%vals(l) = neutrals_parmom_cano(l)**2 / (equi_on_cano%absb(l) * neutrals_dens%vals(l))
            ! Canonical mapping of neutrals parallel velocity
            neutrals_vpar_cano(l) = neutrals_parmom_cano(l) / neutrals_dens%vals(l)
            ! Other canonical mappings for recycling source
            nvpar_cano(l) =  ne%vals(l) * map%flat_stag_to_cano_single(upar%vals, upar%hbwd, l) &
                                        - map%flat_stag_to_cano_single(jpar%vals, jpar%hbwd, l)
            apar_fluct_cano(l) = map%flat_stag_to_cano_single(apar_fluct%vals, apar_fluct%hbwd, l)

        enddo
        !$omp end do
        !$omp do
        do l = 1, mesh_stag%get_n_points()
            neutrals_vpar%vals(l) = neutrals_parmom%vals(l) / map%lift_cano_to_stag_single(neutrals_dens%vals, neutrals_dens%hfwd, l)
            gradpar_neutrals_dens%vals(l) = map%par_grad_single(neutrals_dens%vals, neutrals_dens%hfwd, l)
            gradpar_neutrals_pressure%vals(l) = map%par_grad_single(neutrals_pressure%vals, neutrals_pressure%hfwd, l)
            par_flux_pressure%vals(l) = map%lift_cano_to_stag_single(neutrals_temp%vals, neutrals_temp%hfwd, l) * neutrals_parmom%vals(l)
        enddo
        !$omp end do
        !$omp end parallel
        call perf_stop('../aux_neutrals')

        ! Communicate auxilliary variables --------------------------------------------------------
        call perf_start('../comm_neutrals_2')

        call diff_co%start_comm_halos(comm_handler)
        call diff_co%finalize_comm_halos(comm_handler)

        call gradpar_neutrals_dens%start_comm_halos(comm_handler)
        call gradpar_neutrals_dens%finalize_comm_halos(comm_handler)

        call divbpar_neutrals_parmom%start_comm_halos(comm_handler)
        call divbpar_neutrals_parmom%finalize_comm_halos(comm_handler)

        call par_flux_parmom%start_comm_halos(comm_handler)
        call par_flux_parmom%finalize_comm_halos(comm_handler)

        call neutrals_vpar%start_comm_halos(comm_handler)
        call neutrals_vpar%finalize_comm_halos(comm_handler)

        call par_flux_pressure%start_comm_halos(comm_handler)
        call par_flux_pressure%finalize_comm_halos(comm_handler)

        call gradpar_neutrals_pressure%start_comm_halos(comm_handler)
        call gradpar_neutrals_pressure%finalize_comm_halos(comm_handler)

        call perf_stop('../comm_neutrals_2')

        ! Get penalisation values ----------------------------------------------
        call perf_start('../penalisation')
        call neutrals_dens_penalisation(equi, mesh_cano, map, penalisation_cano, &
                      neutrals_dens, neutrals_dens_pen)
        call neutrals_parmom_penalisation(equi, mesh_stag, map, penalisation_stag, &
                      neutrals_parmom, neutrals_parmom_pen)
        call neutrals_pressure_penalisation(equi, mesh_cano, map, penalisation_cano, &
                      neutrals_dens, neutrals_temp, neutrals_pressure, neutrals_pressure_pen)
        call perf_stop('../penalisation')

        ! Computation of explicit terms ----------------------------
        call perf_start('../main_neutrals')

        ! Compute recycling source
        if ( swn_src_rcy > GP_EPS ) then
            call perf_start('../../recycling_source')
            call compute_recycling_source(comm_handler, equi, mesh_cano, map, equi_on_cano, &
                                          penalisation_cano, parflux_cano, perp_bnd_flux_cano, &
                                          ne, pot, te, nvpar_cano, apar_fluct_cano, &
                                          neutrals_dens, neutrals_parmom_cano, src_rcy)
            call perf_stop('../../recycling_source')
        else
            !$omp parallel do default(none) private(l) shared(mesh_cano, src_rcy)
            do l = 1, mesh_cano%get_n_points()
                src_rcy(l) = 0.0_GP
            enddo
            !$omp end parallel do
        endif

        ! Compute gas puffing/pumping source
        if ( self%gas_puff%source_gaussian_on ) then
            call self%gas_puff%eval_source_gaussian(mesh_cano, map, src_gas_puff)
        else
            !$omp parallel do default(none) private(l) shared(mesh_cano, src_gas_puff)
            do l = 1, mesh_cano%get_n_points()
                src_gas_puff(l) = 0.0_GP
            enddo
            !$omp end parallel do
        endif

        ! Compute diffusion coefficients
        !$omp parallel default(none) private(l) &
        !$omp shared(mesh_cano, mesh_stag, map, diff_co, diff_co_dens, diff_co_pressure, diff_co_parmom, &
        !$omp        swn_dens_diff_pol, swn_pressure_diff_pol, swn_parmom_diff_pol)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            diff_co_dens(l) = swn_dens_diff_pol * diff_co%vals(l)
            diff_co_pressure(l) = swn_pressure_diff_pol * 5.0_GP/3.0_GP * diff_co%vals(l)
        enddo
        !$omp end do
        !$omp do
        do l = 1, mesh_stag%get_n_points()
            diff_co_parmom(l) = swn_parmom_diff_pol * map%lift_cano_to_stag_single(diff_co%vals, diff_co%hfwd, l)
        enddo
        !$omp end do
        !$omp end parallel


        !$omp parallel default(none) &
        !$omp          private(l, ki, kb, kg, x, y, pen_fac, &
        !$omp                  diss_par_neutrals_dens, diff_co_fwdavg, diff_co_bwdavg, diff_tor, &
        !$omp                  diss_par_neutrals_parmom, par_grad_par_flux, par_grad_pressure, &
        !$omp                  diss_par_neutrals_pressure, div_par_flux_pressure, pressure_div_vpar, &
        !$omp                  grad_vpar_sqrd, pressure_vischeat, &
        !$omp                  ddx_neutrals_vpar_cano, ddy_neutrals_vpar_cano) &
        !$omp          shared(equi, tau, mesh_cano, phi_cano, equi_on_cano, mesh_stag, phi_stag, equi_on_stag, &
        !$omp                 penalisation_cano, penalisation_stag, opsinplane_cano, map, &
        !$omp                 k_iz, k_rec, w_iz, p_rec, l_imp,&
        !$omp                 rhos, tratio, pen_epsinv, mms_on, &
        !$omp                 s_cx, pot_iz, te_ref, impy_concentration, rho_min_for_parmom_advection, &
        !$omp                 pardiss_coeff_dens, pardiss_coeff_parmom, pardiss_coeff_pressure, &
        !$omp                 swn_iz, swn_rec, swn_dens_diff_tor, swn_dens_par_flux, &
        !$omp                 swn_parmom_par_flux, swn_parmom_par_pressure, &
        !$omp                 swn_pressure_par_flux, swn_pressure_div_vpar, swn_src_pressure, &
        !$omp                 swn_pressure_vischeat, &
        !$omp                 ne, pot, te, ti, &
        !$omp                 neutrals_dens, neutrals_parmom, neutrals_pressure, neutrals_vpar, neutrals_vpar_cano, &
        !$omp                 gradpar_neutrals_dens, divbpar_neutrals_parmom, par_flux_parmom, &
        !$omp                 gradpar_neutrals_pressure, par_flux_pressure, &
        !$omp                 neutrals_dens_pen, neutrals_parmom_pen, neutrals_pressure_pen, &
        !$omp                 diff_co, src_neutrals_dens, src_neutrals_parmom, src_neutrals_pressure, &
        !$omp                 dneutrals_dens, dneutrals_parmom, dneutrals_pressure, &
        !$omp                 swn_src_rcy, src_rcy, src_gas_puff, src_rcy_temp, tau_adv)

        ! Sources and explicit terms
        !$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)
            pen_fac = (1.0_GP - penalisation_cano%get_charfun_val(ki))

            ! Parallel dynamics
            diss_par_neutrals_dens = map%par_divb_single(gradpar_neutrals_dens%vals, gradpar_neutrals_dens%hbwd, l)

            ! Toroidal diffusion contribution to neutrals particle conservation
            !TODO3D: Toroidal diffusion only works for axisymmetric equilibria
            diff_tor = 0.0_GP
            if (equi%is_axisymmetric() .and. swn_dens_diff_tor > GP_EPS) then
                    diff_co_fwdavg = diff_co%hfwd(l) + diff_co%vals(l)
                    diff_co_bwdavg = diff_co%vals(l) + diff_co%hbwd(l)
                    diff_tor = (diff_co_fwdavg * ( neutrals_pressure%hfwd(l) - neutrals_pressure%vals(l) ) &
                                - diff_co_bwdavg * ( neutrals_pressure%vals(l) - neutrals_pressure%hbwd(l) ) &
                            ) / ( 2.0_GP * ( equi%jacobian(x, y, phi_cano) * map%get_dphi() )**2 )
            endif

            ! Compute explicit part of neutrals density equation
            dneutrals_dens(l) = pen_fac * ( &
                                  + swn_dens_diff_tor * diff_tor & ! explicit toroidal diffusion
                                  - swn_dens_par_flux * divbpar_neutrals_parmom%vals(l) & ! parallel flow
                                  + pardiss_coeff_dens * diss_par_neutrals_dens & ! parallel density dissipation
                                  + src_neutrals_dens(ki) & ! ionization/recombination source
                                  + swn_src_rcy * src_rcy(l) & ! recycling source
                                  + src_gas_puff(l) & ! recycling source
                                ) &
                              + pen_epsinv * penalisation_cano%get_charfun_val(ki) * neutrals_dens_pen(ki) ! penalisation (explicit)

            ! Compute derivatives (scaling due to isotropic normalisation of neutrals in inplane operators)
            ddx_neutrals_vpar_cano = opsinplane_cano%ddx(neutrals_vpar_cano, l) / rhos
            ddy_neutrals_vpar_cano = opsinplane_cano%ddy(neutrals_vpar_cano, l) / rhos

            grad_vpar_sqrd = ddx_neutrals_vpar_cano**2 + ddy_neutrals_vpar_cano**2

            div_par_flux_pressure = map%par_divb_single(par_flux_pressure%vals, par_flux_pressure%hbwd, l)
            pressure_div_vpar = 2.0_GP / 3.0_GP * neutrals_pressure%vals(l) &
                                * map%par_divb_single(neutrals_vpar%vals, neutrals_vpar%hbwd, l)
            pressure_vischeat = 2.0_GP / 3.0_GP * neutrals_pressure%vals(l) * diff_co%vals(l) / tratio * grad_vpar_sqrd
            if (equi_on_cano%rho(l) < rho_min_for_parmom_advection) then
                ! Towards the plasma core, N is close to zero, and computing
                ! parallel neutrals velocity related terms (where a division
                ! by N occurs) can lead to numerical problems.
                div_par_flux_pressure = 0.0_GP
                pressure_div_vpar     = 0.0_GP
                pressure_vischeat     = 0.0_GP
            end if
            diss_par_neutrals_pressure = map%par_divb_single(gradpar_neutrals_pressure%vals, gradpar_neutrals_pressure%hbwd, l)

            ! Compute explicit part of neutrals pressure equation
            dneutrals_pressure(l) = pen_fac * ( &
                                        - swn_pressure_par_flux * div_par_flux_pressure & ! parallel pressure flux
                                        - swn_pressure_div_vpar * pressure_div_vpar & ! parallel compression
                                        + swn_pressure_vischeat * pressure_vischeat & ! viscous heating
                                        + pardiss_coeff_pressure * diss_par_neutrals_pressure & ! parallel pressure dissipation
                                        + swn_src_rcy * src_rcy(l) * src_rcy_temp & ! source from finite temperature of recycled neutrals
                                        + src_neutrals_pressure(ki) & ! neutrals pressure source
                                    ) &
                                    + pen_epsinv * penalisation_cano%get_charfun_val(ki) * neutrals_pressure_pen(ki) ! penalisation (explicit)
        enddo
        !$omp end do
        !$omp do
        do ki = 1, mesh_stag%get_n_points_inner()
            l = mesh_stag%inner_indices(ki)
            x = mesh_stag%get_x(l)
            y = mesh_stag%get_y(l)
            pen_fac = (1.0_GP - penalisation_stag%get_charfun_val(ki))

            ! Parallel dynamics
            diss_par_neutrals_parmom = map%par_grad_single(divbpar_neutrals_parmom%vals, divbpar_neutrals_parmom%hfwd, l)
            par_grad_par_flux = map%par_grad_single(par_flux_parmom%vals, par_flux_parmom%hfwd, l) * equi_on_stag%absb(l)
            if (equi_on_stag%rho(l) < rho_min_for_parmom_advection) then
                ! Towards the plasma core, N is close to zero, and computing parmom
                ! parallel advection as parmom**2/N can lead to numerical problems.
                par_grad_par_flux = 0.0_GP
            end if

            par_grad_pressure = tratio * map%par_grad_single(neutrals_pressure%vals, neutrals_pressure%hfwd, l)

            ! Compute explicit part of neutrals parallel momentum equation
            dneutrals_parmom(l) = pen_fac * ( &
                                    - swn_parmom_par_flux * par_grad_par_flux & ! parallel flow
                                    - swn_parmom_par_pressure * par_grad_pressure & ! par_grad(N*TN)
                                    + pardiss_coeff_parmom * diss_par_neutrals_parmom    & ! parallel dissipation
                                    + src_neutrals_parmom(ki) & ! source
                                  )  &
                                + pen_epsinv * penalisation_stag%get_charfun_val(ki) * neutrals_parmom_pen(ki) ! penalisation (explicit)
        enddo
        !$omp end do

        ! MMS source ----------------------------------------------------------
        if (mms_on) then
            !$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)

                dneutrals_dens(l) = dneutrals_dens(l) &
                                    + mms_source_neutrals_dens(equi, x, y, phi_cano, tau, &
                                                               penalisation_cano%get_charfun_val(ki), &
                                                               pen_epsinv, neutrals_dens_pen(ki), &
                                                               k_iz, k_rec)
                dneutrals_pressure(l) = dneutrals_pressure(l) &
                                        + mms_source_neutrals_pressure(equi, x, y, phi_cano, tau, &
                                                            penalisation_cano%get_charfun_val(ki), &
                                                            pen_epsinv, neutrals_pressure_pen(ki), &
                                                            k_iz, k_rec)
            enddo
            !$omp end do
            !$omp do
            do ki = 1, mesh_stag%get_n_points_inner()
                l = mesh_stag%inner_indices(ki)
                x = mesh_stag%get_x(l)
                y = mesh_stag%get_y(l)

                dneutrals_parmom(l) = dneutrals_parmom(l) &
                                      + mms_source_neutrals_parmom(equi, x, y, phi_stag, tau, &
                                                                   penalisation_stag%get_charfun_val(ki), &
                                                                   pen_epsinv, neutrals_parmom_pen(ki), &
                                                                   k_iz, k_rec)
            enddo
            !$omp end do
        endif

        ! Set ghost points to dummy values for time-step integrators -----------
        !$omp do
        do kb = 1, mesh_cano%get_n_points_boundary()
            l = mesh_cano%boundary_indices(kb)
            dneutrals_dens(l) = 0.0_GP
            dneutrals_pressure(l) = 0.0_GP
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh_cano%get_n_points_ghost()
            l = mesh_cano%ghost_indices(kg)
            dneutrals_dens(l) = 0.0_GP
            dneutrals_pressure(l) = 0.0_GP
        enddo
        !$omp end do
        !$omp do
        do kb = 1, mesh_stag%get_n_points_boundary()
            l = mesh_stag%boundary_indices(kb)
            dneutrals_parmom(l) = 0.0_GP
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh_stag%get_n_points_ghost()
            l = mesh_stag%ghost_indices(kg)
            dneutrals_parmom(l) = 0.0_GP
        enddo
        !$omp end do
        !$omp end parallel
        call perf_stop('../main_neutrals')

        ! Explicit time advancement of variables -------------------------------

        call perf_start('../tadvance_neutrals')
        call tstep_neutrals_dens%advance(dneutrals_dens, neutrals_dens%vals, neutrals_dens_adv)
        call tstep_neutrals_parmom%advance(dneutrals_parmom, neutrals_parmom%vals, neutrals_parmom_adv)
        call tstep_neutrals_pressure%advance(dneutrals_pressure, neutrals_pressure%vals, neutrals_pressure_adv)
        call perf_stop('../tadvance_neutrals')

        ! Get boundary values for MMS analysis ---------------------------------

        if (mms_on) then
            !$omp parallel default(none) private(l, kb, x, y) &
            !$omp          shared(equi, tau_adv, mesh_cano, phi_cano, mesh_stag, phi_stag, &
            !$omp                 boundaries_neutrals)
            !$omp do
             do kb = 1, mesh_cano%get_n_points_boundary()
                l = mesh_cano%boundary_indices(kb)
                x = mesh_cano%get_x(l)
                y = mesh_cano%get_y(l)
                boundaries_neutrals%neutrals_dens%vals(kb) = mms_sol_neutrals_dens(equi, x, y, phi_cano, tau_adv)
                boundaries_neutrals%neutrals_pressure%vals(kb) = mms_sol_neutrals_pressure(equi, x, y, phi_cano, tau_adv)
            enddo
            !$omp end do
            !$omp do
             do kb = 1, mesh_stag%get_n_points_boundary()
                l = mesh_stag%boundary_indices(kb)
                x = mesh_stag%get_x(l)
                y = mesh_stag%get_y(l)
                boundaries_neutrals%neutrals_parmom%vals(kb) = mms_sol_neutrals_parmom(equi, x, y, phi_stag, tau_adv)
            enddo
            !$omp end do
            !$omp end parallel
        endif

        ! Implicit treatment of neutrals density -----------------------------

        call perf_start('../implicit_neutrals')

        call self%setup_implicit_solve(equi, mesh_cano, penalisation_cano, neutrals_dens%vals, &
                        tstep_neutrals_dens, neutrals_dens_adv, boundaries_neutrals%neutrals_dens%vals, &
                        diff_co_dens, neutrals_temp_guess%vals, lambda_cano, xi_cano, co_cano, rhs_cano, &
                        apply_temperature_prefac=.true.)

        if (apply_neumann_mirror) then
            call self%neumann_mirror(equi, mesh_cano, perp_bnd_flux_cano, bnd_descrs_nmn_cano, neutrals_dens%vals, &
                                     tstep_neutrals_dens, diff_co_dens, neutrals_temp_guess%vals, &
                                     lambda_cano, xi_cano, rhs_cano, time_extrapolate=.true., u_floor=floor_dens)
        endif

        call perf_start('../../eslv_neutrals_dens_init')
        call hsolver_cano%update(co_cano, lambda_cano, xi_cano, &
                                 bnddescr_neut_dens_core, &
                                 bnddescr_neut_dens_wall, &
                                 bnddescr_neut_dens_dome, &
                                 bnddescr_neut_dens_out)
        call perf_stop('../../eslv_neutrals_dens_init')

        call perf_start('../../eslv_neutrals_dens_solve')
        call hsolver_cano%solve(rhs_cano, neutrals_dens_adv, res(1), sinfo(1))
        call perf_stop('../../eslv_neutrals_dens_solve')

        ! Implicit treatment of neutrals parmom -----------------------------

        call self%setup_implicit_solve(equi, mesh_stag, penalisation_stag, neutrals_parmom%vals, &
                        tstep_neutrals_parmom, neutrals_parmom_adv, boundaries_neutrals%neutrals_parmom%vals, &
                        diff_co_parmom, neutrals_temp_guess_stag, lambda_stag, xi_stag, co_stag, rhs_stag, &
                        apply_temperature_prefac=.true.)

        if (apply_neumann_mirror) then
            call self%neumann_mirror(equi, mesh_stag, perp_bnd_flux_stag, bnd_descrs_nmn_stag, neutrals_parmom%vals, &
                                     tstep_neutrals_parmom, diff_co_parmom, neutrals_temp_guess_stag, &
                                     lambda_stag, xi_stag, rhs_stag, time_extrapolate=.false.)
        endif

        ! Initialise and perform solve for parallel momentum
        call perf_start('../../eslv_neutrals_parmom_init')
        call hsolver_stag%update(co_stag, lambda_stag, xi_stag, &
                                 bnddescr_neut_parmom_core, &
                                 bnddescr_neut_parmom_wall, &
                                 bnddescr_neut_parmom_dome, &
                                 bnddescr_neut_parmom_out)
        call perf_stop('../../eslv_neutrals_parmom_init')

        call perf_start('../../eslv_neutrals_parmom_solve')
        call hsolver_stag%solve(rhs_stag, neutrals_parmom_adv, res(2), sinfo(2))
        call perf_stop('../../eslv_neutrals_parmom_solve')

        ! Implicit treatment of neutrals pressure -----------------------------

        call self%setup_implicit_solve(equi, mesh_cano, penalisation_cano, neutrals_pressure%vals, &
                        tstep_neutrals_pressure, neutrals_pressure_adv, boundaries_neutrals%neutrals_pressure%vals, &
                        diff_co_pressure, neutrals_temp_guess%vals, lambda_cano, xi_cano, co_cano, rhs_cano, &
                        apply_temperature_prefac=.true.)

        if (apply_neumann_mirror) then
            call self%neumann_mirror(equi, mesh_cano, perp_bnd_flux_cano, bnd_descrs_nmn_cano, neutrals_pressure%vals, &
                                     tstep_neutrals_pressure, diff_co_pressure, neutrals_temp_guess%vals, &
                                     lambda_cano, xi_cano, rhs_cano, time_extrapolate=.false.)
        endif

        ! Initialise and perform solve for pressure
        call perf_start('../../eslv_neutrals_pressure_init')
        call hsolver_cano%update(co_cano, lambda_cano, xi_cano, &
                                 bnddescr_neut_pressure_core, &
                                 bnddescr_neut_pressure_wall, &
                                 bnddescr_neut_pressure_dome, &
                                 bnddescr_neut_pressure_out)
        call perf_stop('../../eslv_neutrals_pressure_init')

        call perf_start('../../eslv_neutrals_pressure_solve')
        call hsolver_cano%solve(rhs_cano, neutrals_pressure_adv, res(3), sinfo(3))
        call perf_stop('../../eslv_neutrals_pressure_solve')

        ! ---- Recover original quantities from implicit solves

        !$omp parallel default(none) private(l) &
        !$omp          shared(mesh_cano, mesh_stag, neutrals_dens_adv, neutrals_parmom_adv, &
        !$omp                 neutrals_pressure_adv, neutrals_temp_guess, neutrals_temp_guess_stag)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            neutrals_dens_adv(l) = neutrals_dens_adv(l) / neutrals_temp_guess%vals(l)
            neutrals_pressure_adv(l) = neutrals_pressure_adv(l) / neutrals_temp_guess%vals(l)
        enddo
        !$omp end do
        !$omp do
        do l = 1, mesh_stag%get_n_points()
            neutrals_parmom_adv(l) = neutrals_parmom_adv(l) / neutrals_temp_guess_stag(l)
        enddo
        !$omp end do
        !$omp end parallel

        ! Set ghost points
        !$omp parallel default(none) private(kg, l, x, y) &
        !$omp          shared(equi, tau_adv, mesh_cano, mesh_stag, phi_cano, phi_stag, mms_on, &
        !$omp                 neutrals_dens_adv, neutrals_parmom_adv, neutrals_pressure_adv)
        if (mms_on) then
            !$omp do
            do kg = 1, mesh_cano%get_n_points_ghost()
                l = mesh_cano%ghost_indices(kg)
                x = mesh_cano%get_x(l)
                y = mesh_cano%get_y(l)
                neutrals_dens_adv(l) = mms_sol_neutrals_dens(equi, x, y, phi_cano, tau_adv)
                neutrals_pressure_adv(l) = mms_sol_neutrals_pressure(equi, x, y, phi_cano, tau_adv)
            enddo
            !$omp end do
            !$omp do
            do kg = 1, mesh_stag%get_n_points_ghost()
                l = mesh_stag%ghost_indices(kg)
                x = mesh_stag%get_x(l)
                y = mesh_stag%get_y(l)
                neutrals_parmom_adv(l) = mms_sol_neutrals_parmom(equi, x, y, phi_stag, tau_adv)
            enddo
            !$omp end do
        else
            call extrapolate_ghost_points(mesh_cano, neutrals_dens_adv)
            call extrapolate_ghost_points(mesh_stag, neutrals_parmom_adv)
            call extrapolate_ghost_points(mesh_cano, neutrals_pressure_adv)
        endif
        !$omp end parallel

        if (evolve_pressure) then
            !$omp parallel do default(none) private(l) &
            !$omp             shared(mesh_cano, neutrals_dens_adv, neutrals_temp, neutrals_pressure_adv)
            do l = 1, mesh_cano%get_n_points()
                neutrals_temp%vals(l) = neutrals_pressure_adv(l) / neutrals_dens_adv(l)
            enddo
            !$omp end parallel do
        endif

        call perf_stop('../implicit_neutrals')

        ! Check for the successful execution of the time step -----------------
        global_sinfo = minval(sinfo)
        call MPI_allreduce(MPI_IN_PLACE, global_sinfo, 1, MPI_Integer, MPI_MIN, &
                                         comm_handler%get_comm_cart(), ierr)
        if ( global_sinfo < 0 ) then
            ! Do not update the data.
            ! Code will safe stop with error snaps output from model_braginskii.
            success_neutrals = .false.
            call perf_stop('timestep_neutrals')
            return
        endif

        success_neutrals = .true.

        ! Formal advancement of time and application of floor-------------------
        call tstep_neutrals_dens%shift_storage( [neutrals_dens%vals, dneutrals_dens] )
        call tstep_neutrals_parmom%shift_storage( [neutrals_parmom%vals, dneutrals_parmom] )
        call tstep_neutrals_pressure%shift_storage( [neutrals_pressure%vals, dneutrals_pressure] )
        call neutrals_temp_store%shift_storage(neutrals_temp%vals)

        ! Reset floor source
        src_floor_neutrals_dens = 0.0_GP

        !$omp parallel default(none) private(l) &
        !$omp          shared(mesh_cano, mesh_stag, floor_dens, floor_temp, &
        !$omp                 neutrals_dens, neutrals_dens_adv, &
        !$omp                 neutrals_parmom, neutrals_parmom_adv, &
        !$omp                 neutrals_pressure, neutrals_pressure_adv, &
        !$omp                 neutrals_temp, evolve_pressure, &
        !$omp                 src_floor_neutrals_dens, dtau)
        !$omp do
        do l = 1, mesh_cano%get_n_points()
            ! Measure floor contribution to density
            src_floor_neutrals_dens(l) = src_floor_neutrals_dens(l) + max(floor_dens - neutrals_dens_adv(l), 0.0_GP) / dtau

            neutrals_dens%vals(l) = max(neutrals_dens_adv(l), floor_dens)
            if (evolve_pressure) then
                neutrals_pressure%vals(l) = max(neutrals_pressure_adv(l), floor_dens*floor_temp)
            else
                neutrals_pressure%vals(l) = neutrals_dens%vals(l) * neutrals_temp%vals(l)
            endif
        enddo
        !$omp end do
        !$omp do
        do l = 1, mesh_stag%get_n_points()
            neutrals_parmom%vals(l) = neutrals_parmom_adv(l)
        enddo
        !$omp end do
        !$omp end parallel

        call perf_stop('timestep_neutrals')

    end subroutine

    subroutine setup_implicit_solve(self, equi, mesh, penalisation, u, tstep_u, u_adv, u_bndval, &
                                    diff_co, temperature_prefac, lambda, xi, co, rhs, apply_temperature_prefac)
        !! Setup equation system for a neutrals diffusion equation on target variable u
        !! by setting lambda, xi, co, rhs coefficients in the helmholtz operator
        class(neutrals_module_t), intent(in) :: self
        !! Instance of class
        class(equilibrium_t), intent(in) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(in) :: mesh
        !! Mesh (any) within poloidal plane
        type(penalisation_t), intent(in) :: penalisation
        !! Penalisation (any)
        real(GP), dimension(mesh%get_n_points()), intent(in) :: u
        !! Target variable values at timepoint t (i.e. before implicit solve)
        type(karniadakis_t), intent(in) :: tstep_u
        !! Timestepper object for target variable
        real(GP), dimension(mesh%get_n_points()), intent(inout) :: u_adv
        !! RHS of implicit solve (e.g. obtained from prior timestepping on explicit components)
        real(GP), dimension(mesh%get_n_points_boundary()), intent(in) :: u_bndval
        !! Boundary values of target variable
        real(GP), dimension(mesh%get_n_points()), intent(in) :: diff_co
        !! Effective diffusion coefficient, excluding the jacobian
        real(GP), dimension(mesh%get_n_points()), intent(in) :: temperature_prefac
        !! Temperature prefactor in lambda when solving for (u T) instead of (u)
        real(GP), dimension(mesh%get_n_points_inner()), intent(out) :: lambda
        !! Output lambda
        real(GP), dimension(mesh%get_n_points_inner()), intent(out) :: xi
        !! Output xi
        real(GP), dimension(mesh%get_n_points()), intent(out) :: co
        !! Output co
        real(GP), dimension(mesh%get_n_points()), intent(out) :: rhs
        !! Output rhs
        logical, intent(in), optional :: apply_temperature_prefac
        !! Switch whether or not to apply temperature prefactor (default: true)
        !! If true, assumes implicit solve is for (u T) instead of (u)
        !! and result of the solve must be divided by T to obtain u at time t+1

        integer :: l, ki, kb, kg
        real(GP) :: x, y, pen_fac, phi
        real(GP), dimension(mesh%get_n_points()) :: temperature_prefac_in
        logical :: apply_temperature_prefac_in

        phi = mesh%get_phi()

        if (present(apply_temperature_prefac)) then
            apply_temperature_prefac_in = apply_temperature_prefac
        else
            apply_temperature_prefac_in = .true.
        endif

        if (apply_temperature_prefac_in) then
            !$omp parallel do default(none) private(l) shared(mesh, temperature_prefac, temperature_prefac_in)
            do l = 1, mesh%get_n_points()
                temperature_prefac_in(l) = temperature_prefac(l)
            enddo
            !$omp end parallel do
        else
            !$omp parallel do default(none) private(l) shared(mesh, temperature_prefac_in)
            do l = 1, mesh%get_n_points()
                temperature_prefac_in(l) = 1.0_GP
            enddo
            !$omp end parallel do
        endif

        !$omp parallel default(none) private(ki, kb, kg, l, x, y, pen_fac) &
        !$omp          shared(equi, mesh, phi, penalisation, pen_epsinv, diff_co, &
        !$omp                 tstep_u, u, u_adv, u_bndval, temperature_prefac_in, &
        !$omp                 lambda, xi, co, rhs)
        !$omp do
        do ki = 1, mesh%get_n_points_inner()
            l = mesh%inner_indices(ki)
            x = mesh%get_x(l)
            y = mesh%get_y(l)
            pen_fac = 1.0_GP - penalisation%get_charfun_val(ki)

            lambda(ki) = ( &
                    1.0_GP &
                    + tstep_u%get_dtau_bdf() * pen_epsinv * penalisation%get_charfun_val(ki) &
                    ) / temperature_prefac_in(l)
            xi(ki) = pen_fac * tstep_u%get_dtau_bdf() / equi%jacobian(x, y, phi)
            co(l) = equi%jacobian(x, y, phi) * diff_co(l)
            rhs(l) = u_adv(l)
            ! Set u_adv to the initial guess for the solver:
            u_adv(l) = temperature_prefac_in(l) * ( &
                                    + 3.0_GP * u(l) &
                                    - 3.0_GP * tstep_u%vstore(l, 1, 1) &
                                    + tstep_u%vstore(l, 2, 1) )
        enddo
        !$omp end do
        !$omp do
        do kb = 1, mesh%get_n_points_boundary()
            l = mesh%boundary_indices(kb)
            x = mesh%get_x(l)
            y = mesh%get_y(l)
            co(l) = equi%jacobian(x, y, phi) * diff_co(l)
            rhs(l) = u_bndval(kb) * temperature_prefac_in(l)
            u_adv(l) = temperature_prefac_in(l) * ( &
                                    + 3.0_GP * u(l) &
                                    - 3.0_GP * tstep_u%vstore(l, 1, 1) &
                                    + tstep_u%vstore(l, 2, 1) )
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh%get_n_points_ghost()
            l = mesh%ghost_indices(kg)
            x = mesh%get_x(l)
            y = mesh%get_y(l)
            co(l) = equi%jacobian(x, y, phi) * diff_co(l)
            ! TODO: With temperature prefactor present, ghost points should really have RHS = u_adv*T
            ! Currently, extrapolate_ghost_pts after the solve fixes this but is not yet consistent.
            rhs(l) = u_adv(l)
            u_adv(l) = temperature_prefac_in(l) * ( &
                                    + 3.0_GP * u(l) &
                                    - 3.0_GP * tstep_u%vstore(l, 1, 1) &
                                    + tstep_u%vstore(l, 2, 1) )
        enddo
        !$omp end do
        !$omp end parallel

    end subroutine

    subroutine neumann_mirror(self, equi, mesh, perp_bnd_flux, bnd_descrs_nmn, u, &
                              tstep_u, diff_co, temperature_prefac, lambda, xi, rhs, &
                              time_extrapolate, u_floor)
        !! Apply "neumann mirroring", i.e. perpendicular neumann zero boundary condition
        !! in preparation for implicit solve of a neutrals equation
        class(neutrals_module_t), intent(in) :: self
        !! Instance of class
        class(equilibrium_t), intent(in) :: equi
        !! Equilibrium
        type(mesh_cart_t), intent(in) :: mesh
        !! Mesh (any) within poloidal plane
        type(perp_bnd_flux_t), intent(in) :: perp_bnd_flux
        !! Perpendicular boundary flux utility
        integer, dimension(mesh%get_n_points_boundary()), intent(in) :: bnd_descrs_nmn
        !! Boundary descriptors
        real(GP), dimension(mesh%get_n_points()), intent(in) :: u
        !! Target variable values at timepoint t (i.e. before implicit solve)
        type(karniadakis_t), intent(in) :: tstep_u
        !! Timestepper object for target variable
        real(GP), dimension(mesh%get_n_points()), intent(in) :: diff_co
        !! Effective diffusion coefficient, excluding the jacobian
        real(GP), dimension(mesh%get_n_points()), intent(in) :: temperature_prefac
        !! Temperature prefactor in lambda when solving for (u T) instead of (u)
        real(GP), dimension(mesh%get_n_points_inner()), intent(inout) :: lambda
        !! Lambda values, some of which will be overwritten at boundary polygon points
        real(GP), dimension(mesh%get_n_points_inner()), intent(inout) :: xi
        !! Xi values, some of which will be overwritten at boundary polygon points
        real(GP), dimension(mesh%get_n_points()), intent(inout) :: rhs
        !! Rhs values, some of which will be overwritten at boundary polygon points
        logical, intent(in) :: time_extrapolate
        !! Switch whether or not to use time extrapolation to mirror values at t+1
        real(GP), intent(in), optional :: u_floor
        !! Optional floor to apply when time extrapolating, only used if time_extrapolate = T

        real(GP), dimension(mesh%get_n_points()) :: wrk
        integer :: is, ks, l, ki, kb, kg

        if ( (cnt > 3) .and. time_extrapolate ) then
            !$omp parallel do default(none) private(l) shared(mesh, wrk, u, tstep_u, u_floor, temperature_prefac)
            do l = 1, mesh%get_n_points()
                wrk(l) = + 3.0_GP * u(l) &
                         - 3.0_GP * tstep_u%vstore(l, 1, 1) &
                         + tstep_u%vstore(l, 2, 1)
                if (present(u_floor)) then
                    wrk(l) = max(wrk(l), u_floor) * temperature_prefac(l)
                else
                    wrk(l) = wrk(l) * temperature_prefac(l)
                endif
            enddo
            !$omp end parallel do
        else
            !$omp parallel do default(none) private(l) shared(mesh, wrk, u, temperature_prefac)
            do l = 1, mesh%get_n_points()
                wrk(l) = u(l) * temperature_prefac(l)
            enddo
            !$omp end parallel do
        endif

        !$omp parallel default(none) private(l, kb, kg) &
        !$omp          shared(mesh, bnd_descrs_nmn, wrk, rhs)
        call set_perpbnds(mesh, bnd_descrs_nmn, wrk)
        call extrapolate_ghost_points(mesh, wrk)
        !$omp do
        do kb = 1, mesh%get_n_points_boundary()
            l = mesh%boundary_indices(kb)
            rhs(l) = wrk(l)
        enddo
        !$omp end do
        !$omp do
        do kg = 1, mesh%get_n_points_ghost()
            l = mesh%ghost_indices(kg)
            rhs(l) = wrk(l)
        enddo
        !$omp end do
        !$omp end parallel

        ! Mirror variable at boundary polygon
        call mirrorset_polyperpbnd_points(equi, mesh, perp_bnd_flux%outer, perp_bnd_flux%conn_outer, &
                                          wrk, diff_co)

        ! Set solve coefficients to leave boundary polygon values unchanged
        !$omp parallel default(none) private(is, ks, l, ki) &
        !$omp          shared(ki_nsegs, mesh, perp_bnd_flux, lambda, xi, rhs, wrk)
        do is = 1, perp_bnd_flux%outer%get_nsegs()
            !$omp do
            do ks = 1, perp_bnd_flux%outer%get_nps(is)
                l = perp_bnd_flux%outer%get_ind_on_mesh(ks, is)
                if (mesh%is_inner_point(l)) then
                    ki = perp_bnd_flux%full_to_inner(l)
                    lambda(ki) = 1.0_GP
                    xi(ki) = 0.0_GP
                endif
                rhs(l) = wrk(l)
            enddo
            !$omp end do
        enddo
        !$omp end parallel

    end subroutine

end module