driver_POST.f90

Go to the documentation of this file.

!------------------------------------------------------------------------------!
!------------------------------------------------------------------------------!
module driver_post
#include <PB3D_macros.h>
#include <wrappers.h>
    use str_utilities
    use output_ops
    use messages
    use num_vars, only: max_str_ln, dp, iu, pi
    use grid_vars, only: grid_type
    use eq_vars, only: eq_1_type, eq_2_type
    use x_vars, only: x_1_type, modes_type, &
        &mds_x, mds_sol
    use sol_vars, only: sol_type
    use vac_vars, only: vac_type
    use rich_vars, only: rich_lvl
    
    implicit none
    private
    public run_driver_post, init_post, stop_post
    
    ! local variables
    integer :: lims_norm(2,3)
    integer, allocatable :: dum_ser(:)
    integer :: rich_lvl_name
    integer :: min_id(3), max_id(3)
    integer :: last_unstable_id
    integer :: n_par_tot
    character(len=4) :: grid_name(3)
    integer :: n_out(3,2)
    type(sol_type) :: sol
    type(vac_type) :: vac
    type(grid_type) :: grid_eq_XYZ
    type(grid_type) :: grid_eq_HEL
    type(grid_type) :: grid_X_HEL
    type(eq_2_type) :: eq_2_HEL
    type(X_1_type) :: X_HEL
    complex(dp), allocatable :: E_pot_int(:,:)
    complex(dp), allocatable :: E_kin_int(:,:)
#if ldebug
    logical :: debug_tor_dep_ripple = .false.                                   
#endif
 
contains
    integer function init_post() result(ierr)
        use grid_vars, only: alpha, n_alpha, min_alpha, max_alpha, n_r_sol
        use num_vars, only: post_style, eq_style, rank, plot_magn_grid, &
            &plot_resonance, plot_flux_q, eq_jobs_lims, plot_grid_style, &
            &n_theta_plot, n_zeta_plot, post_output_full, post_output_sol, &
            &compare_tor_pos, min_r_plot, max_r_plot, min_theta_plot, &
            &max_theta_plot, min_zeta_plot, max_zeta_plot, plot_vac_pot, &
            &x_grid_style, n_procs
        use eq_ops, only: flux_q_plot, divide_eq_jobs, calc_eq_jobs_lims
        use pb3d_ops, only: reconstruct_pb3d_in, reconstruct_pb3d_grid, &
            &reconstruct_pb3d_eq_1, reconstruct_pb3d_eq_2, &
            &reconstruct_pb3d_x_1, reconstruct_pb3d_sol, reconstruct_pb3d_vac, &
            &get_pb3d_grid_size
        use vac_ops, only: vac_pot_plot
        use x_ops, only: init_modes, setup_modes, resonance_plot, calc_res_surf
        use grid_ops, only: calc_norm_range, magn_grid_plot
        use grid_utilities, only: calc_eqd_grid
        use helena_vars, only: nchi
        use sol_ops, only: plot_sol_vals, plot_harmonics
        use mpi_utilities, only: wait_mpi, get_ser_var
        
        character(*), parameter :: rout_name = 'init_POST'
        
        ! local variables
        type(grid_type), target :: grid_eq                                      ! equilibrium grid
        type(grid_type), pointer :: grid_eq_b                                   ! field-aligned equilibrium grid
        type(grid_type) :: grid_x                                               ! perturbation grid
        type(grid_type) :: grid_sol                                             ! solution grid
        type(eq_1_type) :: eq_1                                                 ! flux equilibrium
        integer :: n_div                                                        ! number of divisions of parallel points
        integer :: id, jd                                                       ! counters
        integer :: var_size_without_par(2)                                      ! size without parallel dimension for eq_2 and X_1 variables
        integer :: lim_loc(3,2)                                                 ! grid ranges for local equilibrium job
        integer :: n_div_max                                                    ! maximum divisions of equilibrium jobs
        real(dp), allocatable :: res_surf(:,:)                                  ! resonant surfaces
        real(dp), allocatable :: r_f_x(:)                                       ! normal points in perturbation grid
        character(len=max_str_ln) :: err_msg                                    ! error message
        
        ! initialize ierr
        ierr = 0
        
        ! user output
        select case (post_style)
            case (1)                                                            ! extended grid
                call writo('The post-processing will be done in an extended &
                    &grid')
                call lvl_ud(1)
                call writo('range in r: '//trim(r2strt(min_r_plot))//' .. '//&
                    &trim(r2strt(max_r_plot)))
                call writo('range in θ: '//trim(r2strt(min_theta_plot))//&
                    &' .. '//trim(r2strt(max_theta_plot))//', '//&
                    &trim(i2str(n_theta_plot))//' points')
                call writo('range in ζ: '//trim(r2strt(min_zeta_plot))//&
                    &' .. '//trim(r2strt(max_zeta_plot))//', '//&
                    &trim(i2str(n_zeta_plot))//' points')
                call lvl_ud(-1)
            case (2)                                                            ! field-aligned grid
                call writo('The post-processing will be done in the PB3D grid')
                call lvl_ud(1)
                call writo('range in r: '//trim(r2strt(min_r_plot))//' .. '//&
                    &trim(r2strt(max_r_plot)))
                call lvl_ud(-1)
        end select
        
        ! user output
        call writo('Setting up preliminary variables')
        call lvl_ud(1)
        
        ! some preliminary things
        ierr = reconstruct_pb3d_in('in')                                        ! reconstruct miscellaneous PB3D output variables
        chckerr('')
        
        ! set up alpha
        allocate(alpha(n_alpha))
        ierr = calc_eqd_grid(alpha,min_alpha*pi,max_alpha*pi,excl_last=.true.)
        chckerr('')
        
        ! set up whether Richardson level has to be appended to the name
        select case (eq_style) 
            case (1)                                                            ! VMEC
                rich_lvl_name = rich_lvl                                        ! append richardson level
            case (2)                                                            ! HELENA
                rich_lvl_name = 0                                               ! do not append name
        end select
        
        ! set up grid names
        select case (eq_style)
            case (1)                                                            ! VMEC
                ! take the standard grids
                grid_name(1) = 'eq'
                grid_name(2) = 'X'
                grid_name(3) = 'sol'
            case (2)                                                            ! HELENA
                ! take the field-aligned grids
                grid_name(1) = 'eq_B'
                grid_name(2) = 'X_B'
                grid_name(3) = 'sol'
        end select
        
        ! limits to exclude angular part
        lim_loc(1,:) = [0,0]
        lim_loc(2,:) = [0,0]
        lim_loc(3,:) = [1,-1]
        
        ! reconstruct normal part of full equilibrium, perturbation and solution
        ! grids
        ierr = reconstruct_pb3d_grid(grid_eq,'eq',rich_lvl=rich_lvl_name,&
            &tot_rich=.true.,lim_pos=lim_loc)
        chckerr('')
        ierr = reconstruct_pb3d_grid(grid_x,'X',rich_lvl=rich_lvl_name,&
            &tot_rich=.true.,lim_pos=lim_loc)
        chckerr('')
        ierr = reconstruct_pb3d_grid(grid_sol,'sol')
        chckerr('')
        
        ! set up modes variables in full grids
        !  Note: This  is needed  as many  routines count  on this,  for example
        ! 'set_nm_X' in  X_vars, also  used to  initialize a  solution variable,
        ! etc.
        ierr = reconstruct_pb3d_eq_1(grid_eq,eq_1,'eq_1')
        chckerr('')
        ierr = init_modes(grid_eq,eq_1)
        chckerr('')
        ierr = setup_modes(mds_x,grid_eq,grid_x,plot_name='X_POST')
        chckerr('')
        ierr = setup_modes(mds_sol,grid_eq,grid_sol,plot_name='sol_POST')
        chckerr('')
        
        ! user output
        call lvl_ud(-1)
        call writo('Preliminary variables set up')
        
        ! user output
        call writo('Dividing grids over MPI processes')
        call lvl_ud(1)
        
        ! set eq, X and sol limits, using r_F of the grids
        allocate(r_f_x(grid_x%n(3)))                                            ! r_F_X needs to be allocatable for calc_norm_range
        r_f_x = grid_x%r_F
        ierr = calc_norm_range('POST',eq_limits=lims_norm(:,1),&
            &x_limits=lims_norm(:,2),sol_limits=lims_norm(:,3),&
            &r_f_eq=grid_eq%r_F,r_f_x=r_f_x,r_f_sol=grid_sol%r_F)
        chckerr('')
        deallocate(r_f_x)
        call writo('normal grid limits:')
        call lvl_ud(1)
        call writo('proc '//trim(i2str(rank))//' - equilibrium:  '//&
            &trim(i2str(lims_norm(1,1)))//' .. '//trim(i2str(lims_norm(2,1))),&
            &persistent=.true.)
        call writo('proc '//trim(i2str(rank))//' - perturbation: '//&
            &trim(i2str(lims_norm(1,2)))//' .. '//trim(i2str(lims_norm(2,2))),&
            &persistent=.true.)
        call writo('proc '//trim(i2str(rank))//' - solution:     '//&
            &trim(i2str(lims_norm(1,3)))//' .. '//trim(i2str(lims_norm(2,3))),&
            &persistent=.true.)
        call lvl_ud(-1)
        
        ! synchronize outputs
        ierr = wait_mpi()
        chckerr('')
        
        ! user output
        call lvl_ud(-1)
        call writo('Grids divided over MPI processes')
        
        ! clean up
        call grid_eq%dealloc()
        call grid_x%dealloc()
        call grid_sol%dealloc()
        call eq_1%dealloc()
        
        ! user output
        call writo('Setting up final variables')
        call lvl_ud(1)
        
        ! reconstruct variables
        ierr = reconstruct_pb3d_grid(grid_eq,'eq',rich_lvl=rich_lvl_name,&
            &tot_rich=.true.,grid_limits=lims_norm(:,1))
        chckerr('')
        ierr = reconstruct_pb3d_grid(grid_x,'X',rich_lvl=rich_lvl_name,&
            &tot_rich=.true.,grid_limits=lims_norm(:,2))
        chckerr('')
        ierr = reconstruct_pb3d_grid(grid_sol,'sol',grid_limits=lims_norm(:,3))
        chckerr('')
        ierr = reconstruct_pb3d_eq_1(grid_eq,eq_1,'eq_1')
        chckerr('')
        if (post_output_sol) then
            ierr = reconstruct_pb3d_sol(mds_sol,grid_sol,sol,'sol',&
                &rich_lvl=rich_lvl)
            chckerr('')
            ierr = 2
            chckerr('Vacuum has not been implemented yet!')
            ierr = reconstruct_pb3d_vac(vac,'vac',rich_lvl=rich_lvl_name)
            chckerr('')
        end if
        
        ! user output
        call lvl_ud(-1)
        call writo('Final variables set up')
        
        ! user output
        call writo('1-D PB3D output plots')
        call lvl_ud(1)
        
        if (plot_resonance) then
            ierr = resonance_plot(mds_sol,grid_eq,eq_1)
            chckerr('')
        else
            call writo('Resonance plot not requested')
        end if
        if (plot_flux_q) then
            ierr = flux_q_plot(grid_eq,eq_1)
            chckerr('')
        else
            call writo('Flux quantities plot not requested')
        end if
        if (plot_magn_grid) then
            select case (eq_style)
                case (1)                                                        ! VMEC
                    grid_eq_b => grid_eq
                case (2)                                                        ! HELENA
                    allocate(grid_eq_b)
                    ierr = reconstruct_pb3d_grid(grid_eq_b,'eq_B',&
                        &rich_lvl=rich_lvl,tot_rich=.true.,&
                        &grid_limits=lims_norm(:,1))
                    chckerr('')
            end select
            ierr = magn_grid_plot(grid_eq_b)
            chckerr('')
            if (eq_style.eq.2) then
                call grid_eq_b%dealloc()
                deallocate(grid_eq_b)
            end if
            nullify(grid_eq_b)
        else
            call writo('Magnetic grid plot not requested')
        end if
        
        ierr = wait_mpi()
        chckerr('')
        
        ! user output
        call lvl_ud(-1)
        call writo('1-D outputs plots done')
        
        ! user output
        call writo('Prepare 3-D plots')
        call lvl_ud(1)
        
        ! calculate resonant surfaces on solution grid
        call writo('Calculate resonant surfaces')
        call lvl_ud(1)
        ierr = calc_res_surf(mds_sol,grid_eq,eq_1,res_surf,info=.false.)
        call lvl_ud(-1)
        
        ! plot solution on solution grid
        if (post_output_sol) then
            call writo('Plot the Eigenvalues')
            call lvl_ud(1)
            call plot_sol_vals(sol,last_unstable_id)
            call lvl_ud(-1)
            
            ! find stability regions
            call writo('Find stability ranges')
            call lvl_ud(1)
            call find_stab_ranges(sol,min_id,max_id,last_unstable_id)
            call lvl_ud(-1)
            
            ! plots harmonics for every Eigenvalue, looping over three ranges
            call writo('Plot the Harmonics')
            call lvl_ud(1)
            do jd = 1,3
                if (min_id(jd).le.max_id(jd)) call &
                    &writo('RANGE '//trim(i2str(jd))//': modes '//&
                    &trim(i2str(min_id(jd)))//'..'//trim(i2str(max_id(jd))))
                call lvl_ud(1)
                
                ! indices in each range
                do id = min_id(jd),max_id(jd)
                    ! user output
                    call writo('Mode '//trim(i2str(id))//'/'//&
                        &trim(i2str(size(sol%val)))//', with eigenvalue '&
                        &//trim(c2strt(sol%val(id))))
                    call lvl_ud(1)
                    ierr = plot_harmonics(mds_sol,grid_sol,sol,id,res_surf)
                    chckerr('')
                    
                    ! user output
                    call lvl_ud(-1)
                    call writo('Mode '//trim(i2str(id))//'/'//&
                        &trim(i2str(size(sol%val)))//' finished')
                    
                    ! vacuum plots if requested
                    write(*,*) '¡¡¡¡¡ NO VACUUM !!!!!'
                    plot_vac_pot = .false.
                    if (plot_vac_pot) then
                        ! user output
                        call writo('Start vacuum plot for this solution')
                        call lvl_ud(1)
                        
                        ! plot vacuum potential
                        ierr = vac_pot_plot(grid_sol,vac,sol,id)
                        chckerr('')
                        
                        ! user output
                        call lvl_ud(-1)
                        call writo('Done with vacuum plot')
                    end if
                end do
                
                call lvl_ud(-1)
                if (min_id(jd).le.max_id(jd)) call &
                    &writo('RANGE '//trim(i2str(jd))//' plotted')
            end do
            call lvl_ud(-1)
        end if
        
        ! Divide  the  equilibrium  jobs (see  subroutine  "calc_memory"  inside
        ! "divide_eq_jobs"), only necessary when full output.
        ! The results are  saved in the global variables  eq_jobs_lims for every
        ! eq_job.
        ! If full output  not requested, allocate eq_jobs_lims to  zero size, so
        ! that the driver is never run.
        if (post_output_full) then
            ! set total grid sizes
            do id = 1,2
                ierr = get_pb3d_grid_size(n_out(:,id),grid_name(id),&
                    &rich_lvl=rich_lvl,tot_rich=.true.)
                chckerr('')
                if (post_style.eq.1) n_out(1:2,id) = [n_theta_plot,n_zeta_plot]
            end do
            
            ! test if angular grid sizes are compatible
            if (n_out(1,1).ne.n_out(1,2) .or. n_out(2,1).ne.n_out(2,2)) then
                ierr = 1
                call writo('grid sizes for output grids not compatible. This &
                    &is assumed in the calculation of the equilibrium jobs.')
                call writo('If this has changed, adapt and generalize &
                    &"divide_eq_jobs" appropriately.')
                err_msg = 'Take into account how the ranges transfer from &
                    &grid to grid'
                chckerr(err_msg)
            end if
            
            ! set size of all variables, without parallel dimension
            allocate(dum_ser(2*n_procs))
            ierr = get_ser_var(lims_norm(:,1),dum_ser,scatter=.true.)           ! normal limits of equilibrium grid
            chckerr('')
            var_size_without_par(1) = dum_ser(2*n_procs)-dum_ser(1)+1           ! maximum range
            ierr = get_ser_var(lims_norm(:,2),dum_ser,scatter=.true.)           ! normal limits of perturbation grid
            chckerr('')
            var_size_without_par(2) = dum_ser(2*n_procs)-dum_ser(1)+1           ! maximum range
            select case (x_grid_style)
                case (1,3)                                                      ! equilibrium or enriched
                    var_size_without_par(2) = var_size_without_par(2) + n_r_sol
                case (2)                                                        ! solution
                    ! do nothing
            end select
            deallocate(dum_ser)
            
            n_div_max = n_out(1,1)-1
            if (compare_tor_pos) n_div_max = 1
            select case (eq_style)
                case (1)                                                        ! VMEC
                    ! divide equilibrium jobs
                    ierr = divide_eq_jobs(product(n_out(1:2,1)),&
                        &var_size_without_par,n_div,n_div_max=n_div_max)
                    chckerr('')
                case (2)                                                        ! HELENA
                    ! divide equilibrium jobs
                    ierr = divide_eq_jobs(product(n_out(1:2,1)),&
                        &var_size_without_par,n_div,n_div_max=n_div_max,&
                        &n_par_x_base=nchi)
                    chckerr('')
            end select
            
            ! calculate equilibrium job limits
            ierr = calc_eq_jobs_lims(n_out(1,1),n_div)
            chckerr('')
            
            ! set up total number of parallel points
            n_par_tot = maxval(eq_jobs_lims(2,:))-minval(eq_jobs_lims(1,:))+1
        else
            allocate(eq_jobs_lims(2,0))
        end if
        
        ! Calculate variables  on HELENA output  grid if full output  and HELENA
        ! for later interpolation
        if (post_output_full .and. eq_style.eq.2) then
            ! user output
            call writo('Reconstructing HELENA output for later interpolation')
            call lvl_ud(1)
            
            ierr = reconstruct_pb3d_grid(grid_eq_hel,'eq',&
                &grid_limits=lims_norm(:,1))
            chckerr('')
            ierr = reconstruct_pb3d_grid(grid_x_hel,'X',&
                &grid_limits=lims_norm(:,2))
            chckerr('')
            ierr = reconstruct_pb3d_eq_2(grid_eq_hel,eq_2_hel,'eq_2')
            chckerr('')
            ierr = reconstruct_pb3d_x_1(mds_x,grid_x_hel,x_hel,'X_1')
            chckerr('')
            
            ! user output
            call lvl_ud(-1)
            call writo('HELENA output reconstructed for later interpolation')
        end if
        
        ! reconstruct normal part of full eq grid for XYZ reconstruction
        if (post_output_full .and. plot_grid_style.eq.0 .or. &
            &plot_grid_style.eq.3) then
            ! user output
            call writo('Reconstructing full eq grid for XYZ reconstruction')
            call lvl_ud(1)
            
            ierr = reconstruct_pb3d_grid(grid_eq_xyz,'eq',&
                &rich_lvl=rich_lvl_name,tot_rich=.true.,lim_pos=lim_loc)
            chckerr('')
            
            ! user output
            call lvl_ud(-1)
            call writo('Full eq grid reconstructed for XYZ reconstruction')
        end if
        
        ! user output
        call lvl_ud(-1)
        call writo('3-D plots prepared')
        
        ! clean up
        call grid_eq%dealloc()
        call grid_x%dealloc()
        call grid_sol%dealloc()
        call eq_1%dealloc()
    end function init_post
    
    integer function run_driver_post() result(ierr)
        use num_vars, only: no_output, no_plots, eq_style, eq_jobs_lims, &
            &eq_job_nr, plot_b, plot_j, plot_sol_xi, plot_sol_q, plot_kappa, &
            &post_output_sol, post_style, compare_tor_pos, plot_e_rec
        use pb3d_ops, only: reconstruct_pb3d_grid, reconstruct_pb3d_eq_2, &
            &reconstruct_pb3d_x_1
        use eq_ops, only: calc_eq, calc_t_hf, b_plot, j_plot, &
            &kappa_plot, delta_r_plot
        use eq_utilities, only: calc_inv_met
        use x_ops, only: calc_x
        use sol_ops, only: plot_sol_vec, decompose_energy
        use helena_ops, only: interp_hel_on_grid
        use num_utilities, only: calc_aux_utilities
#if ldebug
        use num_vars, only: ltest
#endif
        
        character(*), parameter :: rout_name = 'run_driver_POST'
        
        ! local variables
        type(grid_type) :: grids(3)                                             ! local output eq, X and sol grid, with limited parallel range and divided normal range
        type(eq_1_type) :: eq_1                                                 ! flux equilibrium
        type(eq_2_type) :: eq_2                                                 ! metric equilibrium
        type(x_1_type) :: x                                                     ! vectorial perturbation variables
        integer :: id, jd                                                       ! counter
        integer :: n_ev_out                                                     ! how many EV's are treated
        integer :: i_ev_out                                                     ! counter
        logical :: no_plots_loc                                                 ! local copy of no_plots
        logical :: no_output_loc                                                ! local copy of no_output
        logical :: last_eq_job                                                  ! last parallel job
        real(dp), allocatable :: xyz_eq(:,:,:,:)                                ! X, Y and Z on equilibrium output grid
        real(dp), allocatable :: xyz_sol(:,:,:,:)                               ! X, Y and Z on solution output grid
        complex(dp), allocatable :: sol_val_comp(:,:,:)                         ! solution eigenvalue for requested solutions
#if ldebug
        logical :: ltest_loc                                                    ! local copy of ltest
#endif
        
        ! initialize ierr
        ierr = 0
        
        ! user output
        call writo('Parallel range for this job:')
        call lvl_ud(1)
        call writo(trim(i2str(eq_jobs_lims(1,eq_job_nr)))//'..'//&
            &trim(i2str(eq_jobs_lims(2,eq_job_nr)))//' of '//'1..'//&
            &trim(i2str(n_par_tot)))
        call lvl_ud(-1)
        
        ! some variables
        last_eq_job = eq_job_nr.eq.size(eq_jobs_lims,2)
        n_ev_out = sum(max_id-min_id+1)
        
        ! set up restricted output grids for this equilibrium job
        ierr = setup_out_grids(grids,xyz_eq,xyz_sol)
        chckerr('')
        
        ! set up output eq_1, eq_2 and X_1, depending on equilibrium style
        ierr = calc_eq(grids(1),eq_1)
        chckerr('')
        select case (eq_style)
            case (1)                                                            ! VMEC
                ! user output
                call writo('Calculate variables on output grid')
                call lvl_ud(1)
                
                ! back up no_plots and no_output and set to .true.
                no_plots_loc = no_plots; no_plots = .true.
                no_output_loc = no_output; no_output = .true.
#if ldebug
                ltest_loc = ltest; ltest = .false.
#endif
                
                ! Calculate the metric equilibrium quantitities
                ierr = calc_eq(grids(1),eq_1,eq_2)
                chckerr('')
                
                ! calculate X variables, vector phase
                ierr = calc_x(mds_x,grids(1),grids(2),eq_1,eq_2,x)
                chckerr('')
                
                ! reset no_plots and no_output
                no_plots = no_plots_loc
                no_output = no_output_loc
#if ldebug
                ltest = ltest_loc
#endif
                
                ! user output
                call lvl_ud(-1)
                call writo('Variables calculated on output grid')
            case (2)                                                            ! HELENA
                ! user output
                call writo('Interpolate variables angularly on output grid')
                call lvl_ud(1)
                
                call eq_2%init(grids(1),setup_e=.true.,setup_f=.true.)
                call x%init(mds_x,grids(2))
                ierr = interp_hel_on_grid(grid_eq_hel,grids(1),&
                    &eq_2=eq_2_hel,eq_2_out=eq_2,eq_1=eq_1,&
                    &grid_name='output equilibrium grid')
                chckerr('')
                ierr = interp_hel_on_grid(grid_x_hel,grids(2),x_1=x_hel,&
                    &x_1_out=x,grid_name='output perturbation grid')
                chckerr('')
                
                ! also set up transformation matrices for plots
                ierr = calc_t_hf(grids(1),eq_1,eq_2,[0,0,0])
                chckerr('')
                ierr = calc_inv_met(eq_2%T_FE,eq_2%T_EF,[0,0,0])
                chckerr('')
                
                ! user output
                call lvl_ud(-1)
                call writo('Variables interpolated angularly on output grid')
        end select
        
        ! user output
        call writo('Initialize 3-D plots')
        call lvl_ud(1)
        
        ! initialize  integrated energies and  open decompose log file  if first
        ! equilibrium job and energy reconstruction requested
        if (plot_e_rec .and. eq_job_nr.eq.1) then
            allocate(e_pot_int(7,n_ev_out))
            allocate(e_kin_int(2,n_ev_out))
            e_pot_int = 0._dp
            e_kin_int = 0._dp
            
            ierr = open_decomp_log()
            chckerr('')
        end if
        
        ! user output
        call lvl_ud(-1)
        call writo('3-D plots initialized')
        
        ! user output
        call writo('3-D PB3D output plots')
        call lvl_ud(1)
        
        ! plot displacement vector if comparing toroidal positions
        if (compare_tor_pos) then
            call writo('Plot the displacement vector')
            call lvl_ud(1)
            ierr = delta_r_plot(grids(1),eq_1,eq_2,xyz_eq)
            chckerr('')
            call lvl_ud(-1)
        end if
        
        ! plot magnetic field if requested
        if (plot_b) then
            call writo('Plot the magnetic field')
            call lvl_ud(1)
            ierr = b_plot(grids(1),eq_1,eq_2,plot_fluxes=post_style.eq.1,&
                &xyz=xyz_eq)
            chckerr('')
            call lvl_ud(-1)
        end if
        
        ! plot current if requested
        if (plot_j) then
            call writo('Plot the current')
            call lvl_ud(1)
            ierr = j_plot(grids(1),eq_1,eq_2,plot_fluxes=post_style.eq.1,&
                &xyz=xyz_eq)
            chckerr('')
            call lvl_ud(-1)
        end if
        
        ! plot magnetic field if requested
        if (plot_kappa) then
            call writo('Plot the curvature')
            call lvl_ud(1)
            ierr = kappa_plot(grids(1),eq_1,eq_2,xyz=xyz_eq)
            chckerr('')
            call lvl_ud(-1)
        end if
        
        if (post_output_sol) then
            ! user output
            call writo('Solution plots for different ranges')
            call lvl_ud(1)
            
            ! loop over three ranges
            i_ev_out = 1
            if (last_eq_job) allocate(sol_val_comp(2,2,n_ev_out))
            do jd = 1,3
                if (min_id(jd).le.max_id(jd)) call &
                    &writo('RANGE '//trim(i2str(jd))//': modes '//&
                    &trim(i2str(min_id(jd)))//'..'//trim(i2str(max_id(jd))))
                call lvl_ud(1)
                
                ! indices in each range
                do id = min_id(jd),max_id(jd)
                    ! user output
                    call writo('Mode '//trim(i2str(id))//'/'//&
                        &trim(i2str(size(sol%val)))//', with eigenvalue '&
                        &//trim(c2strt(sol%val(id))))
                    call lvl_ud(1)
                    
                    ! plot solution vector
                    ierr = plot_sol_vec(mds_x,mds_sol,grids(1),grids(2),&
                        &grids(3),eq_1,eq_2,x,sol,xyz_sol,id,&
                        &[plot_sol_xi,plot_sol_q])
                    chckerr('')
                    
                    if (plot_e_rec) then
                        ! user output
                        call writo('Decompose the energy into its terms')
                        call lvl_ud(1)
                        ierr = decompose_energy(mds_x,mds_sol,grids(1),grids(2),&
                            &grids(3),eq_1,eq_2,x,sol,vac,id,&
                            &b_aligned=post_style.eq.2,xyz=xyz_sol,&
                            &e_pot_int=e_pot_int(:,i_ev_out),&
                            &e_kin_int=e_kin_int(:,i_ev_out))
                        chckerr('')
                        call lvl_ud(-1)
                        
                        ! write  to   log  file  and  save   difference  between
                        ! Eigenvalue and energy fraction if last parallel job
                        if (last_eq_job) then
                            ierr = write_decomp_log(id,e_pot_int(:,i_ev_out),&
                                &e_kin_int(:,i_ev_out))
                            chckerr('')
                            
                            sol_val_comp(:,1,i_ev_out) = id*1._dp
                            sol_val_comp(:,2,i_ev_out) = [sol%val(id),&
                                &sum(e_pot_int(:,i_ev_out))/&
                                &sum(e_kin_int(:,i_ev_out))]
                        end if
                        
                        ! increment counter
                        i_ev_out = i_ev_out + 1
                    end if
                    
                    ! user output
                    call lvl_ud(-1)
                    call writo('Mode '//trim(i2str(id))//'/'//&
                        &trim(i2str(size(sol%val)))//' finished')
                end do
                
                call lvl_ud(-1)
                if (min_id(jd).le.max_id(jd)) call &
                    &writo('RANGE '//trim(i2str(jd))//' plotted')
            end do
            
            ! plot  difference between  Eigenvalue and  energy fraction  if last
            ! parallel job
            if (last_eq_job) then
                call plot_sol_val_comp(sol_val_comp)
            end if
            
            ! user output
            call lvl_ud(-1)
            call writo('Solution for different ranges plotted')
        end if
        
        ! user output
        call lvl_ud(-1)
        call writo('3-D output plots done')
        
        ! clean up
        call writo('Clean up')
        call lvl_ud(1)
        call grids(1)%dealloc()
        call grids(2)%dealloc()
        call grids(3)%dealloc()
        call eq_1%dealloc()
        call eq_2%dealloc()
        call x%dealloc()
        call lvl_ud(-1)
    end function run_driver_post
    
    subroutine stop_post()
        use num_vars, only: plot_grid_style, eq_style, post_output_full, &
            &post_output_sol
        use input_utilities, only: dealloc_in
        
        call dealloc_in()
        if (post_output_sol) then
            call sol%dealloc()
            call vac%dealloc()
        end if
        if (post_output_full .and. eq_style.eq.2) then
            call grid_eq_hel%dealloc()
            call grid_x_hel%dealloc()
            call eq_2_hel%dealloc()
            call x_hel%dealloc()
        end if
        if (post_output_full .and. plot_grid_style.eq.0 .or. &
            &plot_grid_style.eq.3) then
            call grid_eq_xyz%dealloc()
        end if
    end subroutine stop_post
    
    integer function open_decomp_log() result(ierr)
        use num_vars, only: prog_name, output_name, rank, post_style, &
            &decomp_i
        use rich_vars, only: rich_lvl
        
        character(*), parameter :: rout_name = 'open_decomp_log'
        
        ! local variables
        character(len=max_str_ln) :: format_head                                ! header format
        character(len=max_str_ln) :: full_output_name                           ! full name
        character(len=max_str_ln) :: grid_name                                  ! name of grid on which calculations are done
        character(len=2*max_str_ln) :: temp_output_str                          ! temporary output string
        
        ! initialize ierr
        ierr = 0
        
        ! set up output name
        select case (post_style)
            case (1)                                                            ! extended grid
                grid_name = 'extended'
            case (2)                                                            ! field-aligned grid
                grid_name = 'field-aligned'
        end select
        
        call writo('Open decomposition log file')
        call lvl_ud(1)
        if (rank.eq.0) then
            ! set format strings
            format_head = '("#  ",7(A23," "))'
            ! open output file for the log
            full_output_name = prog_name//'_'//trim(output_name)//'_EN_R'//&
                &trim(i2str(rich_lvl))//'.txt'
            open(unit=decomp_i,status='replace',file=full_output_name,&
                &iostat=ierr)
            chckerr('Cannot open EN output file')
            call writo('Log file opened in '//trim(full_output_name))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &'# Energy decomposition using the solution Eigenvectors'
            chckerr('Failed to write')
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &'# (Output on the '//trim(grid_name)//' grid)'
            chckerr('Failed to write')
            write(temp_output_str,format_head) &
                &'RE Eigenvalue          ', 'RE E_pot/E_kin         ', &
                &'RE E_pot               ', 'RE E_kin               '
            write(unit=decomp_i,fmt='(A)',iostat=ierr) trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_head) &
                &'IM Eigenvalue          ', 'IM E_pot/E_kin         ', &
                &'IM E_pot               ', 'IM E_kin               '
            write(unit=decomp_i,fmt='(A)',iostat=ierr) trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_head) &
                &'RE E_kin_n             ', 'RE E_kin_g             '
            write(unit=decomp_i,fmt='(A)',iostat=ierr) trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_head) &
                &'IM E_kin_n             ', 'IM E_kin_g             '
            write(unit=decomp_i,fmt='(A)',iostat=ierr) trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_head) &
                &'RE E_pot line_bending_n', 'RE E_pot line_bending_g', &
                &'RE E_pot ballooning_n  ', 'RE E_pot ballooning_g  ', &
                &'RE E_pot kink_n        ', 'RE E_pot kink_g        ', &
                &'RE E_pot vac           '
            write(unit=decomp_i,fmt='(A)',iostat=ierr) trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_head) &
                &'IM E_pot line_bending_n', 'IM E_pot line_bending_g', &
                &'IM E_pot ballooning_n  ', 'IM E_pot ballooning_g  ', &
                &'IM E_pot kink_n        ', 'IM E_pot kink_g        ', &
                &'IM E_pot vac           '
            write(unit=decomp_i,fmt='(A)',iostat=ierr) trim(temp_output_str)
            chckerr('Failed to write')
        end if
        call lvl_ud(-1)
    end function open_decomp_log
    
    integer function write_decomp_log(X_id,E_pot_int,E_kin_int) result(ierr)
        use num_vars, only: decomp_i, rank
        
        character(*), parameter :: rout_name = 'write_decomp_log'
        
        ! input / output
        integer, intent(in) :: x_id
        complex(dp), intent(in) :: e_pot_int(7)
        complex(dp), intent(in) :: e_kin_int(2)
        
        ! local variables
        character(len=max_str_ln) :: format_val                                 ! format
        character(len=2*max_str_ln) :: temp_output_str                          ! temporary output string
        
        ! initialize ierr
        ierr = 0
        
        ! user output
        call writo('Write to log file')
        call lvl_ud(1)
        
        ! only master
        if (rank.eq.0) then
            ! set up format string:
            !   row 1: EV, E_pot/E_kin
            !   row 2: E_pot, E_kin
            !   row 3: E_kin(1), E_kin(2)
            !   row 4: E_pot(1), E_pot(2)
            !   row 5: E_pot(3), E_pot(4)
            !   row 6: E_pot(5), E_pot(6)
            format_val = '("  ",7(ES23.16," "))'
            
            ! write header
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &'# Eigenvalue '//trim(i2str(x_id))
            chckerr('Failed to write')
            
            ! write values
            write(temp_output_str,format_val) &
                &rp(sol%val(x_id)),&
                &rp(sum(e_pot_int)/sum(e_kin_int)),&
                &rp(sum(e_pot_int)),&
                &rp(sum(e_kin_int))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_val) &
                &ip(sol%val(x_id)),&
                &ip(sum(e_pot_int)/sum(e_kin_int)), &
                &ip(sum(e_pot_int)),&
                &ip(sum(e_kin_int))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_val) &
                &rp(e_kin_int(1)),&
                &rp(e_kin_int(2))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_val) &
                &ip(e_kin_int(1)),&
                &ip(e_kin_int(2))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_val) &
                &rp(e_pot_int(1)),&
                &rp(e_pot_int(2)),&
                &rp(e_pot_int(3)),&
                &rp(e_pot_int(4)),&
                &rp(e_pot_int(5)),&
                &rp(e_pot_int(6)),&
                &rp(e_pot_int(7))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &trim(temp_output_str)
            chckerr('Failed to write')
            write(temp_output_str,format_val) &
                &ip(e_pot_int(1)),&
                &ip(e_pot_int(2)),&
                &ip(e_pot_int(3)),&
                &ip(e_pot_int(4)),&
                &ip(e_pot_int(5)),&
                &ip(e_pot_int(6)),&
                &ip(e_pot_int(7))
            write(unit=decomp_i,fmt='(A)',iostat=ierr) &
                &trim(temp_output_str)
            chckerr('Failed to write')
            
            close(decomp_i)
        end if
        
        call lvl_ud(-1)
    end function write_decomp_log
    
    subroutine find_stab_ranges(sol,min_id,max_id,last_unstable_id)
        use num_vars, only: n_sol_plotted
        
        ! input / output
        type(sol_type), intent(in) :: sol
        integer, intent(inout) :: min_id(3)
        integer, intent(inout) :: max_id(3)
        integer, intent(inout) :: last_unstable_id
        
        ! local variables
        integer :: id                                                           ! counter
        integer :: n_sol_found                                                  ! how many solutions found and saved
        
        ! set local variables
        n_sol_found = size(sol%val)
        
        ! find last unstable index (if ends with 0, no unstable EV)
        last_unstable_id = 0
        do id = 1,n_sol_found
            if (rp(sol%val(id)).lt.0._dp) last_unstable_id = id
        end do
        ! set up min. and max. of range 1
        if (last_unstable_id.gt.0) then                                         ! there is an unstable range
            if (n_sol_plotted(1).ge.0) then
                min_id(1) = 1                                                   ! start from most unstable EV
                max_id(1) = min(n_sol_plotted(1),last_unstable_id)              ! end with n_sol_plotted(1) first unstable values if available
            else
                min_id(1) = 1                                                   ! start from most unstable EV
                max_id(1) = last_unstable_id                                    ! end with last unstable EV
            end if
        else                                                                    ! no unstable range
            min_id(1) = 1                                                       ! no unstable values to plot
            max_id(1) = 0                                                       ! no unstable values to plot
        end if
        ! set up min. and max. of range 2
        if (last_unstable_id.gt.0) then                                         ! there is an unstable range
            if (n_sol_plotted(2).ge.0) then
                min_id(2) = last_unstable_id - n_sol_plotted(2) + 1             ! start from n_sol_plotted(2) last unstable values
            else
                min_id(2) = 1                                                   ! start from 1
            end if
            if (n_sol_plotted(3).ge.0) then
                max_id(2) = min(last_unstable_id + n_sol_plotted(3),n_sol_found)! end with n_sol_plotted(3) first stable values if available
            else
                max_id(2) = n_sol_found                                         ! end with last EV
            end if
        else                                                                    ! no unstable range
            min_id(2) = 1                                                       ! start from first EV (stable)
            if (n_sol_plotted(3).ge.0) then
                max_id(2) = min(n_sol_plotted(3),n_sol_found)                   ! end with n_sol_plotted(3) first stable values if available
            else
                max_id(2) = n_sol_found                                         ! end with last EV
            end if
        end if
        ! set up min. and max. of range 3
        if (n_sol_plotted(4).ge.0) then
            min_id(3) = n_sol_found - n_sol_plotted(4) + 1                      ! start from n_sol_plotted(4) last stable values
        else
            min_id(3) = last_unstable_id + 1                                    ! start from n_sol_plotted(4) last stable values
        end if
        max_id(3) = n_sol_found                                                 ! end with most stable EV
        ! merge ranges 2 and 3 if overlap
        if (min_id(3).le.max_id(2)) then
            max_id(2) = max_id(3)                                               ! range 3 merged with range 2
            min_id(3) = 1                                                       ! range 3 not used any more
            max_id(3) = 0                                                       ! range 3 not used any more
        end if
        ! merge ranges 1 and 2 if overlap
        if (min_id(2).le.max_id(1)) then
            max_id(1) = max_id(2)                                               ! range 2 merged with range 1
            min_id(2) = 1                                                       ! range 2 not used any more
            max_id(2) = 0                                                       ! range 2 not used any more
        end if
    end subroutine find_stab_ranges
    
    subroutine plot_sol_val_comp(sol_val_comp)
        use num_vars, only: rank
        
        ! input /output
        complex(dp), intent(inout) :: sol_val_comp(:,:,:)
        
        ! local variables
        character(len=max_str_ln) :: plot_title(2)                              ! title for plots
        character(len=max_str_ln) :: plot_name                                  ! file name for plots
        
        if (rank.eq.0 .and. size(sol_val_comp,3).gt.0) then
            ! real part
            plot_title = ['RE sol_val','E_frac    ']
            plot_name = 'sol_val_comp_RE'
            call print_ex_2d(plot_title,plot_name,&
                &rp(sol_val_comp(:,1,:)),&
                &x=rp(sol_val_comp(:,2,:)),draw=.false.)
            call draw_ex(plot_title,plot_name,size(sol_val_comp,3),1,.false.)
            plot_title(1) = 'rel diff between RE sol_val and E_frac'
            plot_name = 'sol_val_comp_RE_rel_diff'
            call print_ex_2d(plot_title(1),plot_name,rp(&
                &(sol_val_comp(1,2,:)-sol_val_comp(2,2,:))/&
                &sol_val_comp(1,2,:)),x=rp(sol_val_comp(1,1,:)),&
                &draw=.false.)
            call draw_ex(plot_title(1:1),plot_name,1,1,.false.)
            plot_title(1) = 'rel diff between RE sol_val and E_frac [log]'
            plot_name = 'sol_val_comp_RE_rel_diff_log'
            call print_ex_2d(plot_title(1),plot_name,log10(abs(rp(&
                &(sol_val_comp(1,2,:)-sol_val_comp(2,2,:))/&
                &sol_val_comp(1,2,:)))),x=rp(sol_val_comp(1,1,:)),&
                &draw=.false.)
            call draw_ex(plot_title(1:1),plot_name,1,1,.false.)
            
            ! imaginary part
            plot_title = ['IM sol_val','E_frac    ']
            plot_name = 'sol_val_comp_IM'
            call print_ex_2d(plot_title,plot_name,&
                &rp(sol_val_comp(:,1,:)),x=ip(sol_val_comp(:,2,:)),draw=.false.)
            call draw_ex(plot_title,plot_name,size(sol_val_comp,3),1,.false.)
            plot_title(1) = 'rel diff between IM sol_val and E_frac'
            plot_name = 'sol_val_comp_IM_rel_diff'
            call print_ex_2d(plot_title(1),plot_name,ip(&
                &(sol_val_comp(1,2,:)-sol_val_comp(2,2,:))/&
                &sol_val_comp(1,2,:)),x=rp(sol_val_comp(1,1,:)),&
                &draw=.false.)
            call draw_ex(plot_title(1:1),plot_name,1,1,.false.)
            plot_title(1) = 'rel diff between IM sol_val and E_frac [log]'
            plot_name = 'sol_val_comp_IM_rel_diff_log'
            call print_ex_2d(plot_title(1),plot_name,log10(abs(ip(&
                &(sol_val_comp(1,2,:)-sol_val_comp(2,2,:))/&
                &sol_val_comp(1,2,:)))),x=rp(sol_val_comp(1,1,:)),&
                &draw=.false.)
            call draw_ex(plot_title(1:1),plot_name,1,1,.false.)
        end if
    end subroutine plot_sol_val_comp
    
    integer function setup_out_grids(grids_out,XYZ_eq,XYZ_sol) result(ierr)
        use num_vars, only: min_theta_plot, max_theta_plot, post_style, &
            &eq_job_nr, eq_jobs_lims, norm_disc_prec_x, x_grid_style
        use num_utilities, only: spline
        use grid_utilities, only: extend_grid_f
        use pb3d_ops, only: reconstruct_pb3d_grid
        use num_vars, only: rank, n_procs
#if ldebug
        use grid_utilities, only: nufft
#endif
        
        character(*), parameter :: rout_name = 'setup_out_grids'
        
        ! input / output
        type(grid_type), intent(inout) :: grids_out(3)
        real(dp), intent(inout), allocatable :: xyz_eq(:,:,:,:)
        real(dp), intent(inout), allocatable :: xyz_sol(:,:,:,:)
        
        ! local variables
        type(grid_type) :: grid_eq                                              ! equilibrium grid
        type(grid_type) :: grid_x                                               ! perturbation grid
        integer :: n_theta                                                      ! nr. of points in theta for extended grid
        integer :: id, jd, ld                                                   ! counters
        integer :: lim_loc(3,2,3)                                               ! grid ranges for local equilibrium job (last index: eq, X, sol)
        real(dp) :: lim_theta(2)                                                ! theta limits
        real(dp), allocatable :: r_geo(:,:,:)                                   ! geometrical radius
        real(dp), allocatable :: xy(:,:,:)                                      ! x and y
        real(dp), allocatable :: f(:,:,:)                                       ! Fourier components
        real(dp), allocatable :: f_loc(:,:)                                     ! local f
        real(dp), allocatable :: xyz_x(:,:,:,:)                                 ! X, Y and Z on output perturbation grid
        
        ! initialize ierr
        ierr = 0
        
        call writo('Set up output grid')
        call lvl_ud(1)
        
        select case (post_style)
            case (1)                                                            ! extended grid
                ! limits to exclude angular part
                call writo('Set up limits')
                ! eq
                lim_loc(1,:,1) = [0,0]
                lim_loc(2,:,1) = [0,0]
                lim_loc(3,:,1) = [1,-1]
                ! X
                lim_loc(1,:,2) = [0,0]
                lim_loc(2,:,2) = [0,0]
                lim_loc(3,:,2) = [1,-1]
                ! sol
                lim_loc(1,:,3) = [0,0]
                lim_loc(2,:,3) = [0,0]
                lim_loc(3,:,3) = [1,-1]
                
                ! reconstruct equilibrium, perturbation and solution grids
                call writo('Reconstruct PB3D variables')
                ierr = reconstruct_pb3d_grid(grid_eq,'eq',&
                    &rich_lvl=rich_lvl_name,lim_pos=lim_loc(:,:,1),&
                    &grid_limits=lims_norm(:,1))
                chckerr('')
                ierr = reconstruct_pb3d_grid(grid_x,'X',&
                    &rich_lvl=rich_lvl_name,lim_pos=lim_loc(:,:,2),&
                    &grid_limits=lims_norm(:,2))
                chckerr('')
                ierr = reconstruct_pb3d_grid(grids_out(3),'sol',&
                    &lim_pos=lim_loc(:,:,3),grid_limits=lims_norm(:,3))
                chckerr('')
                
                ! theta limits
                if (n_par_tot.gt.1) then
                    lim_theta = min_theta_plot + &
                        &(eq_jobs_lims(:,eq_job_nr)-1._dp)*&
                        &(max_theta_plot-min_theta_plot)/(n_par_tot-1._dp)
                else
                    lim_theta = min_theta_plot
                end if
                n_theta = eq_jobs_lims(2,eq_job_nr)-eq_jobs_lims(1,eq_job_nr)+1
                
                ! extend eq and X grid
                call writo('Extend equilibrium grid')
                call lvl_ud(1)
                ierr = extend_grid_f(grid_eq,grids_out(1),n_theta_plot=n_theta,&
                    &lim_theta_plot=lim_theta,grid_eq=grid_eq)                  ! extend eq grid and convert to F
                chckerr('')
                call lvl_ud(-1)
                
                call writo('Extend perturbation grid')
                call lvl_ud(1)
                ierr = grids_out(2)%init([n_theta,n_out(2:3,2)],&
                    &i_lim=[grid_x%i_min,grid_x%i_max])
                chckerr('')
                grids_out(2)%r_F = grid_x%r_F
                grids_out(2)%r_E = grid_x%r_E
                grids_out(2)%loc_r_F = grid_x%loc_r_F
                grids_out(2)%loc_r_E = grid_x%loc_r_E
                select case (x_grid_style)
                    case (1)                                                    ! equilibrium
                        ! copy angular variables
                        grids_out(2)%theta_F = grids_out(1)%theta_F
                        grids_out(2)%theta_E = grids_out(1)%theta_E
                        grids_out(2)%zeta_F = grids_out(1)%zeta_F
                        grids_out(2)%zeta_E = grids_out(1)%zeta_E
                    case (2,3)                                                  ! solution or enriched
                        ! interpolate angular variables
                        do jd = 1,grids_out(1)%n(2)
                            do id = 1,grids_out(1)%n(1)
                                ierr = spline(grids_out(1)%loc_r_F,&
                                    &grids_out(1)%theta_F(id,jd,:),&
                                    &grids_out(2)%loc_r_F,&
                                    &grids_out(2)%theta_F(id,jd,:),&
                                    &ord=norm_disc_prec_x)
                                chckerr('')
                                ierr = spline(grids_out(1)%loc_r_F,&
                                    &grids_out(1)%theta_E(id,jd,:),&
                                    &grids_out(2)%loc_r_F,&
                                    &grids_out(2)%theta_E(id,jd,:),&
                                    &ord=norm_disc_prec_x)
                                chckerr('')
                                ierr = spline(grids_out(1)%loc_r_F,&
                                    &grids_out(1)%zeta_F(id,jd,:),&
                                    &grids_out(2)%loc_r_F,&
                                    &grids_out(2)%zeta_F(id,jd,:),&
                                    &ord=norm_disc_prec_x)
                                chckerr('')
                                ierr = spline(grids_out(1)%loc_r_F,&
                                    &grids_out(1)%zeta_E(id,jd,:),&
                                    &grids_out(2)%loc_r_F,&
                                    &grids_out(2)%zeta_E(id,jd,:),&
                                    &ord=norm_disc_prec_x)
                                chckerr('')
                            end do
                        end do
                end select
                call lvl_ud(-1)
                
                ! clean up
                call grid_eq%dealloc()
                call grid_x%dealloc()
            case (2)                                                            ! field-aligned grid
                ! limits to take subset of angular part
                call writo('Set up limits')
                ! eq
                lim_loc(1,:,1) = eq_jobs_lims(:,eq_job_nr)
                lim_loc(2,:,1) = [1,-1]
                lim_loc(3,:,1) = [1,-1]
                ! X
                lim_loc(1,:,2) = eq_jobs_lims(:,eq_job_nr)
                lim_loc(2,:,2) = [1,-1]
                lim_loc(3,:,2) = [1,-1]
                ! sol
                lim_loc(1,:,3) = [0,0]
                lim_loc(2,:,3) = [0,0]
                lim_loc(3,:,3) = [1,-1]
                
                ! reconstruct equilibrium, perturbation and solution grids
                call writo('Reconstruct PB3D variables')
                ierr = reconstruct_pb3d_grid(grids_out(1),trim(grid_name(1)),&
                    &rich_lvl=rich_lvl,tot_rich=.true.,&
                    &lim_pos=lim_loc(:,:,1),grid_limits=lims_norm(:,1))
                chckerr('')
                ierr = reconstruct_pb3d_grid(grids_out(2),trim(grid_name(2)),&
                    &rich_lvl=rich_lvl,tot_rich=.true.,&
                    &lim_pos=lim_loc(:,:,2),grid_limits=lims_norm(:,2))
                chckerr('')
                ierr = reconstruct_pb3d_grid(grids_out(3),trim(grid_name(3)),&
                    &lim_pos=lim_loc(:,:,3),grid_limits=lims_norm(:,3))
                chckerr('')
        end select
        
        call writo('Calculate X, Y and Z of equilibrium grid')
        call lvl_ud(1)
        ierr = calc_xyz_of_output_grid(grids_out(1),xyz_eq)
        chckerr('')
        call lvl_ud(-1)
        
        call writo('Calculate X, Y and Z of perturbation grid')
        call lvl_ud(1)
        ierr = calc_xyz_of_output_grid(grids_out(2),xyz_x)
        chckerr('')
        call lvl_ud(-1)
        
        call writo('Calculate X, Y and Z of solution grid')
        call lvl_ud(1)
        allocate(xyz_sol(grids_out(2)%n(1),grids_out(2)%n(2),&
            &grids_out(3)%loc_n_r,3))
        select case (x_grid_style)
            case (1,3)                                                          ! equilibrium or enriched
                ! setup normal interpolation data
                do ld = 1,3
                    do jd = 1,grids_out(2)%n(2)
                        do id = 1,grids_out(2)%n(1)
                            ierr = spline(grids_out(2)%loc_r_F,&
                                &xyz_x(id,jd,:,ld),grids_out(3)%loc_r_F,&
                                &xyz_sol(id,jd,:,ld),ord=norm_disc_prec_x)
                            chckerr('')
                        end do
                    end do
                end do
            case (2)                                                            ! solution
                xyz_sol = xyz_x
        end select
        call lvl_ud(-1)
        
#if ldebug
        ! Plot toroidal dependency of ripple.
        ! Note: Hard-code NFP.
        if (debug_tor_dep_ripple .and. rank.eq.n_procs-1) then
            call plot_hdf5(['X','Y','Z'],'XYZ',xyz_eq)
            allocate(r_geo(grids_out(1)%n(1),grids_out(1)%n(2),&
                &grids_out(1)%loc_n_r))
            r_geo = sqrt((xyz_eq(:,:,:,3)-0.56710)**2 + &
                &(xyz_eq(:,:,:,1)-6.4297)**2)
            call plot_hdf5('r_geo','r_geo_TEMP',r_geo,x=xyz_eq(:,:,:,1),&
                &y=xyz_eq(:,:,:,2),z=xyz_eq(:,:,:,3))
            allocate(xy(grids_out(1)%n(2),grids_out(1)%n(1),2))
            do id = 1,grids_out(1)%n(1)
                xy(:,id,1) = xyz_eq(id,:,grids_out(1)%loc_n_r,2)*18
                xy(:,id,2) = r_geo(id,:,grids_out(1)%loc_n_r) &
                    &* grids_out(1)%n(2) / sum(r_geo(id,:,grids_out(1)%loc_n_r))
            end do
            call print_ex_2d(['r_geo'],'r_geo',xy(:,:,2),x=xy(:,:,1),&
                &draw=.false.)
            call draw_ex(['r_geo'],'r_geo',grids_out(1)%n(1),1,.false.,&
                &ex_plot_style=1)
            do id = 1,grids_out(1)%n(1)
                ierr = nufft(xy(1:size(xy,1)-1,id,1),xy(1:size(xy,1)-1,id,2),&
                    &f_loc)
                chckerr('')
                if (id.eq.1) &
                    &allocate(f(size(f_loc,1),size(f_loc,2),grids_out(1)%n(1)))
                f(:,:,id) = f_loc
            end do
            call print_ex_2d(['r_geo_cos'],'r_geo_cos',f(:,1,:),&
                &draw=.false.)
            call draw_ex(['r_geo_cos'],'r_geo_cos',grids_out(1)%n(1),1,.false.,&
                &ex_plot_style=1)
            call print_ex_2d(['r_geo_sin'],'r_geo_sin',f(:,2,:),&
                &draw=.false.)
            call draw_ex(['r_geo_sin'],'r_geo_sin',grids_out(1)%n(1),1,.false.,&
                &ex_plot_style=1)
        end if
#endif
        
        call lvl_ud(-1)
    contains
        ! Calculate XYZ of output grid, depending on plot_grid style:
        !   0: 3-D geometry
        !   1: slab geometry
        !   2: slab geometry with wrapping of angles
        !   3: 3-D geometry with straightened toroidal coordinate
        integer function calc_xyz_of_output_grid(grid,XYZ) result(ierr)
            use grid_utilities, only: calc_xyz_grid
            use num_vars, only: eq_style, plot_grid_style, swap_angles, &
                &use_pol_flux_f
            use eq_vars, only: max_flux_f
            use grid_vars, only: alpha, n_alpha
            use vmec_utilities, onLy: calc_trigon_factors
            use mpi_utilities, only: wait_mpi
            
            character(*), parameter :: rout_name = 'calc_XYZ_of_output_grid'
            
            ! input / output
            type(grid_type), intent(inout) :: grid                              ! output grid for which to calculate XYZ
            real(dp), allocatable :: xyz(:,:,:,:)                               ! X, Y and Z on output grid
            
            ! local variables
            real(dp), allocatable :: r(:,:,:)                                   ! R on output grid
            character(len=max_str_ln) :: err_msg                                ! error message
            integer :: jd, kd                                                   ! counters
            
            ! initialize ierr
            ierr = 0
            
            ! if VMEC, calculate trigonometric factors of output grid
            if (eq_style.eq.1) then
                ierr = calc_trigon_factors(grid%theta_E,grid%zeta_E,&
                    &grid%trigon_factors)
                chckerr('')
            end if
            
            ! set up XYZ
            allocate(xyz(grid%n(1),grid%n(2),grid%loc_n_r,3))
            select case (plot_grid_style)
                case (0)                                                        ! 3-D geometry
                    ! user output
                    call writo('Plots are done in 3-D geometry')
                    call lvl_ud(1)
                    
                    ierr = calc_xyz_grid(grid_eq_xyz,grid,&
                        &xyz(:,:,:,1),xyz(:,:,:,2),xyz(:,:,:,3))
                    chckerr('')
                    !!! To plot the cross-section
                    !!call print_ex_2D(['cross_section'],'cross_section_'//&
                        !!&trim(i2str(eq_job_nr))//'_'//trim(i2str(rank)),&
                        !!&XYZ(:,1,:,3),x=XYZ(:,1,:,1),draw=.false.)
                    !!call draw_ex(['cross_section'],'cross_section_'//&
                        !!&trim(i2str(eq_job_nr))//'_'//trim(i2str(rank)),&
                        !!&size(XYZ,3),1,.false.)
                    
                    call lvl_ud(-1)
                case (3)                                                        ! 3-D geometry with straightened toroidal coordinate
                    ! user output
                    call writo('Plots are done in an unwrapped torus')
                    call lvl_ud(1)
                    
                    allocate(r(grid%n(1),grid%n(2),grid%loc_n_r))
                    ierr = calc_xyz_grid(grid_eq_xyz,grid,&
                        &xyz(:,:,:,1),xyz(:,:,:,2),xyz(:,:,:,3),r=r)
                    chckerr('')
                    xyz(:,:,:,1) = r                                            ! set X to R coordinate
                    xyz(:,:,:,2) = grid%zeta_F                                  ! set Y to toroidal coordinate
                    !!! To plot the cross-section
                    !!call print_ex_2D(['cross_section'],'cross_section_'//&
                        !!&trim(i2str(eq_job_nr))//'_'//trim(i2str(rank)),&
                        !!&XYZ(:,1,:,3),x=XYZ(:,1,:,1),draw=.false.)
                    !!call draw_ex(['cross_section'],'cross_section_'//&
                        !!&trim(i2str(eq_job_nr))//'_'//trim(i2str(rank)),&
                        !!&size(XYZ,3),1,.false.)
                    
                    call lvl_ud(-1)
                case (1,2)                                                      ! slab geometry (with or without wrapping to fundamental interval)
                    ! user output
                    call writo('Plots are done in slab geometry')
                    call lvl_ud(1)
                    
                    ! set angular coordinates
                    select case (post_style)
                        case (1)                                                ! extended grid: both theta and zeta need to be chosen, not alpha
                            if (use_pol_flux_f) then
                                xyz(:,:,:,1) = grid%theta_F/pi
                                xyz(:,:,:,2) = grid%zeta_F/pi
                            else
                                xyz(:,:,:,1) = grid%zeta_F/pi
                                xyz(:,:,:,2) = grid%theta_F/pi
                            end if
                        case (2)                                                ! field-aligned grid: alpha in combination with parallel angle
                            if (swap_angles) then
                                call writo('For plotting, the angular &
                                    &coordinates are swapped: theta <-> zeta')
                                if (use_pol_flux_f) then
                                    xyz(:,:,:,1) = grid%zeta_F/pi
                                else
                                    xyz(:,:,:,1) = grid%theta_F/pi
                                end if
                            else
                                if (use_pol_flux_f) then
                                    xyz(:,:,:,1) = grid%theta_F/pi
                                else
                                    xyz(:,:,:,1) = grid%zeta_F/pi
                                end if
                            end if
                            do jd = 1,n_alpha
                                xyz(:,jd,:,2) = alpha(jd)
                            end do
                    end select
                    
                    ! set normal coordinate
                    do kd = 1,grid%loc_n_r
                        xyz(:,:,kd,3) = grid%loc_r_F(kd)/max_flux_f*2*pi
                    end do
                    
                    ! limit to fundamental interval -1..1 if plot grid style 2
                    if (plot_grid_style.eq.2) xyz(:,:,:,1:2) = &
                        &modulo(xyz(:,:,:,1:2)+1._dp,2._dp)-1._dp
                    
                    call lvl_ud(-1)
                case default
                    err_msg = 'No style "'//trim(i2str(plot_grid_style))//&
                        &'" for plot_grid_style'
                    ierr = 1
                    chckerr(err_msg)
            end select
            
            ! synchronize
            ierr = wait_mpi()
            chckerr('')
        end function calc_xyz_of_output_grid
    end function setup_out_grids
end module driver_post