grid_ops.f90

Go to the documentation of this file.

!------------------------------------------------------------------------------!
!------------------------------------------------------------------------------!
module grid_ops
#include <PB3D_macros.h>
    use str_utilities
    use output_ops
    use messages
    use num_vars, only: dp, pi, max_str_ln
    use grid_vars, only: grid_type
    use eq_vars, only: eq_1_type
 
    implicit none
    private
    public calc_norm_range, calc_ang_grid_eq_b, magn_grid_plot, setup_grid_eq, &
        &setup_grid_eq_b, setup_grid_x, setup_grid_sol, print_output_grid, &
        &redistribute_output_grid
#if ldebug
    public debug_calc_ang_grid_eq_b
#endif
    
    ! global variables
#if ldebug
 
    logical :: debug_calc_ang_grid_eq_b = .false.                               
#endif
 
contains
    integer function calc_norm_range(style,in_limits,eq_limits,X_limits,&
        &sol_limits,r_F_eq,r_F_X,r_F_sol,jq) result(ierr)
        
        character(*), parameter :: rout_name = 'calc_norm_range'
        
        ! input / output
        character(len=*), intent(in) :: style
        integer, intent(inout), optional :: in_limits(2)
        integer, intent(inout), optional :: eq_limits(2)
        integer, intent(inout), optional :: x_limits(2)
        integer, intent(inout), optional :: sol_limits(2)
        real(dp), intent(inout), optional :: r_f_eq(:)
        real(dp), intent(inout), allocatable, optional :: r_f_x(:)
        real(dp), intent(inout), optional :: r_f_sol(:)
        real(dp), intent(in), optional :: jq(:)
        
        ! local variables
        character(len=max_str_ln) :: err_msg                                    ! error message
        
        ! initialize ierr
        ierr = 0
        
        ! select depending on style
        select case (style)
            case ('PB3D_in')                                                    ! PB3D: in
                if (present(in_limits)) then
                    ierr = calc_norm_range_pb3d_in(in_limits)
                    chckerr('')
                else
                    ierr = 1
                    err_msg = 'for PB3D: in, in_limits need to be provided'
                    chckerr(err_msg)
                end if
            case ('PB3D_eq')                                                    ! PB3D: eq
                if (present(eq_limits)) then
                    call calc_norm_range_pb3d_eq(eq_limits)
                else
                    ierr = 1
                    err_msg = 'for PB3D: eq, eq_limits need to be provided'
                    chckerr(err_msg)
                end if
            case ('PB3D_X')                                                     ! PB3D: X
                if (present(eq_limits).and.present(x_limits).and.&
                    &present(r_f_eq).and.present(r_f_x).and.(present(jq))) then
                    ierr = calc_norm_range_pb3d_x(eq_limits,x_limits,&
                        &r_f_eq,r_f_x,jq)
                    chckerr('')
                else
                    ierr = 1
                    err_msg = 'for PB3D: X, eq_limits, X_limits, r_F_eq, &
                        &r_F_X and jq need to be provided'
                    chckerr(err_msg)
                end if
            case ('PB3D_sol')                                                   ! PB3D: sol
                if (present(sol_limits).and.present(r_f_sol)) then
                    ierr = calc_norm_range_pb3d_sol(sol_limits,r_f_sol)
                    chckerr('')
                else
                    ierr = 1
                    err_msg = 'for PB3D: sol, sol_limits and r_F_sol need to &
                        &be provided'
                    chckerr(err_msg)
                end if
            case ('POST')                                                       ! POST
                if (present(eq_limits) .and. present(x_limits) .and. &
                    &present(sol_limits) .and.  present(r_f_eq) .and. &
                    &present(r_f_x) .and.  present(r_f_sol)) then
                    call calc_norm_range_post(eq_limits,x_limits,sol_limits,&
                        &r_f_eq,r_f_x,r_f_sol)
                else
                    ierr = 1
                    err_msg = 'for POST, eq_limits, X_limits, sol_limits, &
                        &r_F_eq, r_F_X and r_F_sol need to be provided'
                    chckerr(err_msg)
                end if
            case default
                ierr = 1
                err_msg = 'Incorrect style "'//trim(style)//'"'
                chckerr(err_msg)
        end select
    contains
        integer function calc_norm_range_pb3d_in(in_limits) &
            &result(ierr)                                                       ! PB3D version for equilibrium grid
            use num_vars, only: use_pol_flux_e, use_pol_flux_f, eq_style, &
                &norm_disc_prec_eq
            use num_utilities, only: con2dis, dis2con, round_with_tol, spline
            use vmec_vars, only: flux_p_v, flux_t_v
            use helena_vars, only: flux_p_h, flux_t_h
            use x_vars, only: min_r_sol, max_r_sol
            use eq_vars, only: max_flux_e, max_flux_f
            use grid_vars, only: n_r_in
            
            character(*), parameter :: rout_name = 'calc_norm_range_PB3D_in'
            
            ! input / output
            integer, intent(inout) :: in_limits(2)
            
            ! local variables
            real(dp), allocatable :: flux_f(:), flux_e(:)                       ! either pol. or tor. flux in F and E
            real(dp) :: tot_min_r_in_f_con                                      ! tot_min_r_in in continuous F coords.
            real(dp) :: tot_max_r_in_f_con                                      ! tot_max_r_in in continuous F coords.
            real(dp) :: tot_min_r_in_e_con(1)                                   ! tot_min_r_in in continuous E coords.
            real(dp) :: tot_max_r_in_e_con(1)                                   ! tot_max_r_in in continuous E coords.
            real(dp) :: tot_min_r_in_e_dis                                      ! tot_min_r_in in discrete E coords., unrounded
            real(dp) :: tot_max_r_in_e_dis                                      ! tot_max_r_in in discrete E coords., unrounded
            
            ! initialize ierr
            ierr = 0
            
            ! set up flux_F and flux_E,  depending on which equilibrium style is
            ! being used:
            !   1:  VMEC
            !   2:  HELENA
            allocate(flux_f(n_r_in),flux_e(n_r_in))
            select case (eq_style)
                case (1)                                                        ! VMEC
                    ! set up F flux
                    if (use_pol_flux_f) then
                        flux_f = flux_p_v(:,0)
                    else
                        flux_f = - flux_t_v(:,0)                                ! conversion VMEC LH -> RH coord. system
                    end if
                    ! set up E flux
                    if (use_pol_flux_e) then
                        flux_f = flux_p_v(:,0)
                    else
                        flux_e = flux_t_v(:,0)
                    end if
                case (2)                                                        ! HELENA
                    ! set up F flux
                    if (use_pol_flux_f) then
                        flux_f = flux_p_h(:,0)
                    else
                        flux_f = flux_t_h(:,0)
                    end if
                    ! set up E flux
                    if (use_pol_flux_e) then
                        flux_e = flux_p_h(:,0)
                    else
                        flux_e = flux_t_h(:,0)
                    end if
            end select
            
            ! set max_flux
            max_flux_e = flux_e(n_r_in)
            max_flux_f = flux_f(n_r_in)
            
            ! normalize flux_F and flux_E to (0..1) using max_flux
            flux_e = flux_e/max_flux_e
            flux_f = flux_f/max_flux_f
            
            ! get lower and upper bound of total solution range
            ! 1. start
            tot_min_r_in_f_con = min_r_sol
            tot_max_r_in_f_con = max_r_sol
            ! 2. round with tolerance
            ierr = round_with_tol(tot_min_r_in_f_con,0._dp,1._dp)
            chckerr('')
            ierr = round_with_tol(tot_max_r_in_f_con,0._dp,1._dp)
            chckerr('')
            ! 3. continuous E coords
            ierr = spline(flux_f,flux_e,[tot_min_r_in_f_con],&
                &tot_min_r_in_e_con,ord=norm_disc_prec_eq)
            ierr = spline(flux_f,flux_e,[tot_max_r_in_f_con],&
                &tot_max_r_in_e_con,ord=norm_disc_prec_eq)
            ! 4. round with tolerance
            ierr = round_with_tol(tot_min_r_in_e_con,minval(flux_e),&
                &maxval(flux_e))
            chckerr('')
            ierr = round_with_tol(tot_max_r_in_e_con,minval(flux_e),&
                &maxval(flux_e))
            chckerr('')
            ! 5. discrete E index, unrounded
            ierr = con2dis(tot_min_r_in_e_con(1),tot_min_r_in_e_dis,flux_e)
            chckerr('')
            ierr = con2dis(tot_max_r_in_e_con(1),tot_max_r_in_e_dis,flux_e)
            chckerr('')
            ! 6. discrete E index, rounded
            in_limits(1) = floor(tot_min_r_in_e_dis)
            in_limits(2) = ceiling(tot_max_r_in_e_dis)
        end function calc_norm_range_pb3d_in
        
        subroutine calc_norm_range_pb3d_eq(eq_limits)                           ! PB3D version for equilibrium grid
            use num_vars, only: n_procs, rank, norm_disc_prec_eq
            use grid_vars, only: n_r_eq
            
            ! input / output
            integer, intent(inout) :: eq_limits(2)
            
            ! set local equilibrium limits
            eq_limits(1) = nint(1 + 1._dp*rank*(n_r_eq-1)/n_procs)
            eq_limits(2) = nint(1 + (1._dp+rank)*(n_r_eq-1)/n_procs)
            
            ! increase lower limits for processes that are not first
            if (rank.gt.0) eq_limits(1) = eq_limits(1)+1
            
            ! ghost regions of width 4*norm_disc_prec_eq
            ! Note: When finite differences were used, a factor of 2 was enough,
            ! but it  looks like for  splines it has to  be higher. Even  for 4,
            ! there is some slight discrepancy...  This should be looked into in
            ! the future.
            eq_limits(1) = max(eq_limits(1)-4*norm_disc_prec_eq,1)
            eq_limits(2) = min(eq_limits(2)+4*norm_disc_prec_eq,n_r_eq)
        end subroutine calc_norm_range_pb3d_eq
        
        integer function calc_norm_range_pb3d_x(eq_limits,X_limits,r_F_eq,&
            &r_F_X,jq) result(ierr)                                             ! PB3D version for perturbation grid
            
            use num_vars, only: x_grid_style, max_njq_change, rank
            use grid_vars, only: n_r_eq, n_r_sol
            use x_vars, only: prim_x
            use grid_utilities, only: calc_eqd_grid
            use eq_vars, only: max_flux_f
            
            character(*), parameter :: rout_name = 'calc_norm_range_PB3D_X'
            
            ! input / output
            integer, intent(in) :: eq_limits(2)
            integer, intent(inout) :: x_limits(2)
            real(dp), intent(in) :: r_f_eq(:)
            real(dp), intent(inout), allocatable :: r_f_x(:)
            real(dp), intent(in) :: jq(:)
            
            ! local variables
            integer :: kd                                                       ! counter
            integer, allocatable :: div(:)                                      ! number of extra divisions for this grid interval
            real(dp), allocatable :: r_f_plot(:,:)                              ! plot of r_F_eq and r_F_X
            
            ! initialize ierr
            ierr = 0
            
            select case (x_grid_style)
                case (1)                                                        ! equilibrium
                    ! copy equilibrium
                    allocate(r_f_x(n_r_eq))
                    r_f_x = r_f_eq
                    x_limits = eq_limits
                    
                    ! output
                    call writo('Using the equilibrium grid')
                case (2)                                                        ! solution
                    ! calculate solution
                    allocate(r_f_x(n_r_sol))
                    ierr = calc_norm_range_pb3d_sol(sol_limits=x_limits,&
                        &r_f_sol=r_f_x)
                    chckerr('')
                    
                    ! output
                    call writo('Using the solution grid')
                case (3)                                                        ! enriched
                    ! decide extra divisions for each interval
                    allocate(div(size(r_f_eq)-1))
                    do kd = 1,size(r_f_eq)-1
                        div(kd) = &
                            &floor(prim_x*abs(jq(kd+1)-jq(kd))/max_njq_change)
                    end do
                    
                    ! set up r_F_X with divisions
                    allocate(r_f_x(size(r_f_eq)+sum(div)))
                    do kd = 1,size(r_f_eq)-1
                        ierr = calc_eqd_grid(&
                            &r_f_x(kd+sum(div(1:kd-1)):kd+sum(div(1:kd))),&
                            &r_f_eq(kd),r_f_eq(kd+1),excl_last=.true.)
                        chckerr('')
                    end do
                    r_f_x(size(r_f_x)) = r_f_eq(size(r_f_eq))
                    
                    ! set up normal limits
                    x_limits(1) = eq_limits(1)+sum(div(1:eq_limits(1)-1))
                    x_limits(2) = eq_limits(2)+sum(div(1:eq_limits(2)-1))
                    
                    ! plot
                    if (rank.eq.0) then
                        allocate(r_f_plot(size(r_f_x),4))
                        r_f_plot(1:size(r_f_eq),1) = r_f_eq
                        r_f_plot(1:size(r_f_eq),3) = [(kd,kd=1,size(r_f_eq))]
                        r_f_plot(size(r_f_eq)+1:size(r_f_x),1) = &
                            &r_f_eq(size(r_f_eq))
                        r_f_plot(size(r_f_eq)+1:size(r_f_x),3) = size(r_f_eq)
                        r_f_plot(:,2) = r_f_x
                        r_f_plot(:,4) = [(kd,kd=1,size(r_f_x))]
                        r_f_plot(:,1:2) = r_f_plot(:,1:2)/max_flux_f*2*pi
                        call print_ex_2d(['eq','X '],'r_F_X',&
                            &r_f_plot(:,3:4),x=r_f_plot(:,1:2),draw=.false.)
                        call draw_ex(['eq','X '],'r_F_X',2,1,&
                            &.false.)
                    end if
                    
                    ! output
                    call writo('Enriching the equilibrium grid to have jq &
                        &change not more than '//trim(r2strt(max_njq_change)))
                    call lvl_ud(1)
                    if (any(div.gt.0)) then
                        call writo('Added '//trim(i2str(sum(div)))//' points')
                    else
                        call writo('But it was not necessary to  add points')
                    end if
                    call lvl_ud(-1)
            end select
        end function calc_norm_range_pb3d_x
        
        integer function calc_norm_range_pb3d_sol(sol_limits,r_F_sol,n_procs) &
            &result(ierr)                                                       ! PB3D version for solution grid
            use num_vars, only: sol_n_procs, rank, norm_disc_prec_sol
            use eq_vars, only: max_flux_f
            use x_vars, only: min_r_sol, max_r_sol
            use num_utilities, only: round_with_tol
            use grid_vars, only: n_r_sol
            
            character(*), parameter :: rout_name = 'calc_norm_range_PB3D_sol'
            
            ! input / output
            integer, intent(inout) :: sol_limits(2)
            real(dp), intent(inout) :: r_f_sol(:)
            integer, intent(in), optional :: n_procs
            
            ! local variables
            integer :: id                                                       ! counter
            integer, allocatable :: loc_n_r_sol(:)                              ! local nr. of normal points in solution grid
            integer :: n_procs_loc                                              ! local n_procs
            
            ! initialize ierr
            ierr = 0
            
            ! set local n_procs
            n_procs_loc = sol_n_procs
            if (present(n_procs)) n_procs_loc = n_procs
            
            ! set loc_n_r_sol
            allocate(loc_n_r_sol(n_procs_loc))
            loc_n_r_sol = n_r_sol/n_procs_loc                                   ! number of radial points on this processor
            loc_n_r_sol(1:mod(n_r_sol,n_procs_loc)) = &
                &loc_n_r_sol(1:mod(n_r_sol,n_procs_loc)) + 1                    ! add a mode to if there is a remainder
            
            ! set sol_limits
            if (rank.lt.n_procs_loc) then
                sol_limits = &
                    &[sum(loc_n_r_sol(1:rank))+1,sum(loc_n_r_sol(1:rank+1))]
                
                ! ghost regions of width 2*norm_disc_prec_sol
                sol_limits(1) = max(sol_limits(1)-norm_disc_prec_sol,1)
                sol_limits(2) = min(sol_limits(2)+norm_disc_prec_sol,n_r_sol)
            else
                sol_limits(1) = n_r_sol
                sol_limits(2) = n_r_sol
            end if
            
            ! calculate r_F_sol in range from min_r_sol to max_r_sol
            r_f_sol = [(min_r_sol + (id-1.0_dp)/(n_r_sol-1.0_dp)*&
                &(max_r_sol-min_r_sol),id=1,n_r_sol)]
            
            ! round with standard tolerance
            ierr = round_with_tol(r_f_sol,0.0_dp,1.0_dp)
            chckerr('')
            
            ! translate  to   the  real   normal  variable  in   range  from
            ! 0..flux/2pi
            r_f_sol = r_f_sol*max_flux_f/(2*pi)
        end function calc_norm_range_pb3d_sol
        
        subroutine calc_norm_range_post(eq_limits,X_limits,sol_limits,r_F_eq,&
            &r_F_X,r_F_sol)                                                     ! POST version
            
            use num_vars, only: n_procs, rank, norm_disc_prec_sol, &
                &min_r_plot, max_r_plot, x_grid_style
            use eq_vars, only: max_flux_f
            use grid_utilities, only: find_compr_range
            
            ! input / output
            integer, intent(inout) :: eq_limits(2)
            integer, intent(inout) :: X_limits(2)
            integer, intent(inout) :: sol_limits(2)
            real(dp), intent(in) :: r_F_eq(:)
            real(dp), intent(in) :: r_F_X(:)
            real(dp), intent(in) :: r_F_sol(:)
            
            ! local variables
            integer :: sol_limits_tot(2)                                        ! total solution limits
            integer :: n_r_sol                                                  ! total nr. of normal points in solution grid
            integer, allocatable :: loc_n_r_sol(:)                              ! local nr. of normal points in solution grid
            real(dp) :: min_sol, max_sol                                        ! min. and max. of r_F_sol in range of this process
            
            ! initialize ierr
            ierr = 0
            
            ! get min and max of solution range
            min_sol = max(minval(r_f_sol),min_r_plot*max_flux_f/(2*pi))
            max_sol = min(maxval(r_f_sol),max_r_plot*max_flux_f/(2*pi))
            
            ! find the solution index that comprises this range
            call find_compr_range(r_f_sol,[min_sol,max_sol],sol_limits_tot)
            
            ! initialize n_r_sol
            n_r_sol = sol_limits_tot(2)-sol_limits_tot(1)+1
            allocate(loc_n_r_sol(n_procs))
            
            ! divide the solution grid equally over all the processes
            loc_n_r_sol = n_r_sol/n_procs                                       ! number of radial points on this processor
            loc_n_r_sol(1:mod(n_r_sol,n_procs)) = &
                &loc_n_r_sol(1:mod(n_r_sol,n_procs)) + 1                        ! add a mode to if there is a remainder
            
            ! set sol_limits
            sol_limits = [sum(loc_n_r_sol(1:rank))+1,sum(loc_n_r_sol(1:rank+1))]
            sol_limits = sol_limits + sol_limits_tot(1) - 1
            if (rank.gt.0) sol_limits(1) = sol_limits(1)-norm_disc_prec_sol     ! ghost region for num. deriv.
            if (rank.lt.n_procs-1) sol_limits(2) = &
                &sol_limits(2)+norm_disc_prec_sol+1                             ! ghost region for num. deriv. and overlap for int_vol
            min_sol = minval(r_f_sol(sol_limits(1):sol_limits(2)))
            max_sol = maxval(r_f_sol(sol_limits(1):sol_limits(2)))
            
            ! determine eq_limits: smallest eq and X range comprising sol range
            call find_compr_range(r_f_eq,[min_sol,max_sol],eq_limits)
            select case (x_grid_style)
                case (1,3)                                                      ! equilibrium or enriched
                    call find_compr_range(r_f_x,[min_sol,max_sol],x_limits)
                case (2)                                                        ! solution
                    x_limits = sol_limits
            end select
        end subroutine calc_norm_range_post
    end function calc_norm_range
 
    integer function setup_grid_eq(grid_eq,eq_limits) result(ierr)
        use num_vars, only: eq_style, eq_jobs_lims, eq_job_nr
        use grid_vars, only: n_r_eq, n_alpha
        use helena_vars, only: nchi, chi_h
        use rich_vars, only: n_par_x
        
        character(*), parameter :: rout_name = 'setup_grid_eq'
        
        ! input / output
        type(grid_type), intent(inout) :: grid_eq
        integer, intent(in) :: eq_limits(2)
        
        ! local variables
        integer :: id                                                           ! counter
        
        ! initialize ierr
        ierr = 0
        
        ! set up general equilibrium grid:
        ! choose which equilibrium style is being used:
        !   1:  VMEC
        !   2:  HELENA
        select case (eq_style)
            case (1)                                                            ! VMEC
                ! user output
                call writo('Field-aligned with '//trim(i2str(n_r_eq))//&
                    &' normal and '//trim(i2str(n_par_x))//' parallel points')
                
                ! create grid with eq_jobs_lims
                ierr = grid_eq%init([eq_jobs_lims(2,eq_job_nr)-&
                    &eq_jobs_lims(1,eq_job_nr)+1,n_alpha,n_r_eq],eq_limits)
                chckerr('')
            case (2)                                                            ! HELENA
                ! user output
                call writo('Identical to the equilibrium input grid')
                call writo(trim(i2str(n_r_eq))//' normal and '//&
                    &trim(i2str(nchi))//' angular points')
                call writo('Poloidal range '//trim(r2strt(chi_h(1)))//'..'//&
                    &trim(r2strt(chi_h(nchi))))
                call writo('Will be interpolated to a field-aligned grid &
                    &later')
                
                ! create grid
                ierr = grid_eq%init([nchi,1,n_r_eq],eq_limits)                  ! axisymmetric equilibrium
                chckerr('')
                
                ! copy angular grid from HELENA
                do id = 1,grid_eq%n(1)
                    grid_eq%theta_E(id,:,:) = chi_h(id)
                end do
                grid_eq%zeta_E = 0._dp
                
                ! convert to Flux coordinates (trivial)
                grid_eq%theta_F = grid_eq%theta_E
                grid_eq%zeta_F = grid_eq%zeta_E
        end select
    end function setup_grid_eq
    
    integer function setup_grid_eq_b(grid_eq,grid_eq_B,eq,only_half_grid) &
        &result(ierr)
        use num_vars, only: eq_style, eq_jobs_lims, eq_job_nr
        use grid_vars, only: n_alpha
        
        character(*), parameter :: rout_name = 'setup_grid_eq_B'
        
        ! input / output
        type(grid_type), intent(inout) :: grid_eq
        type(grid_type), intent(inout) :: grid_eq_b
        type(eq_1_type), intent(in) :: eq
        logical, intent(in), optional :: only_half_grid
        
        ! local variables
        character(len=max_str_ln) :: err_msg                                    ! error message
        
        ! initialize ierr
        ierr = 0
        
        ! choose which equilibrium style is being used:
        !   1:  VMEC
        !   2:  HELENA
        select case (eq_style)
            case (1)                                                            ! VMEC
                ! the grid is already field aligned
                ierr = 1
                err_msg = 'The grid is already field-aligned for VMEC'
                chckerr(err_msg)
            case (2)                                                            ! HELENA
                ! create grid with eq_jobs_lims
                ierr = grid_eq_b%init([eq_jobs_lims(2,eq_job_nr)-&
                    &eq_jobs_lims(1,eq_job_nr)+1,n_alpha,grid_eq%n(3)],&
                    &[grid_eq%i_min,grid_eq%i_max])                             ! only one field line
                chckerr('')
                
                ! copy the normal coords.
                grid_eq_b%loc_r_E = grid_eq%loc_r_E
                grid_eq_b%loc_r_F = grid_eq%loc_r_F
                grid_eq_b%r_E = grid_eq%r_E
                grid_eq_b%r_F = grid_eq%r_F
                
                ! calculate the angular grid that follows the magnetic field
                ierr = calc_ang_grid_eq_b(grid_eq_b,eq,only_half_grid)
                chckerr('')
        end select
    end function setup_grid_eq_b
    
    integer function setup_grid_x(grid_eq,grid_X,r_F_X,X_limits) result(ierr)
        use num_vars, only: norm_disc_prec_x, x_grid_style
        use grid_utilities, only: coord_f2e
        use num_utilities, only: spline
        
        character(*), parameter :: rout_name = 'setup_grid_X'
        
        ! input / output
        type(grid_type), intent(in) :: grid_eq
        type(grid_type), intent(inout) :: grid_x
        real(dp), intent(in), allocatable :: r_f_x(:)
        integer, intent(in) :: x_limits(2)
        
        ! local variables
        integer :: id, jd                                                       ! counters
        
        ! initialize ierr
        ierr = 0
        
        select case (x_grid_style)
            case (1)                                                            ! equilibrium
                ! X grid identical to equilibrium grid
                ierr = grid_eq%copy(grid_x)
                chckerr('')
            case (2,3)                                                          ! solution and enriched
                ! create grid
                ierr = grid_x%init([grid_eq%n(1:2),size(r_f_x)],x_limits)
                chckerr('')
                
                ! set Flux variables
                grid_x%r_F = r_f_x
                grid_x%loc_r_F = r_f_x(x_limits(1):x_limits(2))
                
                ! convert to Equilibrium variables
                ierr = coord_f2e(grid_eq,grid_x%r_F,grid_x%r_E,&
                    &r_f_array=grid_eq%r_F,r_e_array=grid_eq%r_E)
                chckerr('')
                ierr = coord_f2e(grid_eq,grid_x%loc_r_F,grid_x%loc_r_E,&
                    &r_f_array=grid_eq%r_F,r_e_array=grid_eq%r_E)
                chckerr('')
                
                ! interpolate
                do jd = 1,grid_eq%n(2)
                    do id = 1,grid_eq%n(1)
                        ierr = spline(grid_eq%loc_r_F,grid_eq%theta_E(id,jd,:),&
                            &grid_x%loc_r_F,grid_x%theta_E(id,jd,:),&
                            &ord=norm_disc_prec_x)
                        chckerr('')
                        ierr = spline(grid_eq%loc_r_F,grid_eq%zeta_E(id,jd,:),&
                            &grid_x%loc_r_F,grid_x%zeta_E(id,jd,:),&
                            &ord=norm_disc_prec_x)
                        chckerr('')
                        ierr = spline(grid_eq%loc_r_F,grid_eq%theta_F(id,jd,:),&
                            &grid_x%loc_r_F,grid_x%theta_F(id,jd,:),&
                            &ord=norm_disc_prec_x)
                        chckerr('')
                        ierr = spline(grid_eq%loc_r_F,grid_eq%zeta_F(id,jd,:),&
                            &grid_x%loc_r_F,grid_x%zeta_F(id,jd,:),&
                            &ord=norm_disc_prec_x)
                        chckerr('')
                    end do
                end do
        end select
    end function setup_grid_x
    
    integer function setup_grid_sol(grid_eq,grid_X,grid_sol,r_F_sol,&
        &sol_limits) result(ierr)
        
        use num_vars, only: n_procs, x_grid_style
        use grid_utilities, only: coord_f2e
        
        character(*), parameter :: rout_name = 'setup_grid_sol'
        
        ! input / output
        type(grid_type), intent(in) :: grid_eq
        type(grid_type), intent(in) :: grid_x
        type(grid_type), intent(inout) :: grid_sol
        real(dp), intent(in) :: r_f_sol(:)
        integer, intent(in) :: sol_limits(2)
        
        ! initialize ierr
        ierr = 0
        
        ! output
        call writo(trim(i2str(size(r_f_sol)))//' normal  points')
        
        select case (x_grid_style)
            case (1,3)                                                          ! equilibrium or enriched
                ! create grid
                ierr = grid_sol%init([0,0,size(r_f_sol)],sol_limits,&
                    &divided=n_procs.gt.1)
                chckerr('')
                
                ! set Flux variables
                grid_sol%r_F = r_f_sol
                grid_sol%loc_r_F = r_f_sol(sol_limits(1):sol_limits(2))
                
                ! convert to Equilibrium variables
                ierr = coord_f2e(grid_eq,grid_sol%r_F,grid_sol%r_E,&
                    &r_f_array=grid_eq%r_F,r_e_array=grid_eq%r_E)
                chckerr('')
                ierr = coord_f2e(grid_eq,grid_sol%loc_r_F,grid_sol%loc_r_E,&
                    &r_f_array=grid_eq%r_F,r_e_array=grid_eq%r_E)
                chckerr('')
            case (2)                                                            ! solution
                ! solution grid  identical to perturation  grid but with  only 1
                ! parallel point.
                ierr = grid_sol%init([0,0,grid_x%n(3)],&
                    &[grid_x%i_min,grid_x%i_max],grid_x%divided)
                chckerr('')
                grid_sol%r_F = grid_x%r_F
                grid_sol%loc_r_F = grid_x%r_F(sol_limits(1):sol_limits(2))
                grid_sol%r_E = grid_x%r_E
                grid_sol%loc_r_E = grid_x%r_E(sol_limits(1):sol_limits(2))
        end select
    end function setup_grid_sol
    
    integer function calc_ang_grid_eq_b(grid_eq,eq,only_half_grid) result(ierr)
        use num_vars, only: use_pol_flux_f, use_pol_flux_e, &
            &eq_style, tol_zero, eq_job_nr, eq_jobs_lims
        use grid_vars, only: min_par_x, max_par_x, alpha, n_alpha
        use eq_vars, only: max_flux_e
        use grid_utilities, only: coord_f2e, calc_eqd_grid, calc_n_par_x_rich
        use eq_utilities, only: print_info_eq
        use rich_vars, only: n_par_x
        use x_vars, only: min_r_sol, max_r_sol
#if ldebug
        use grid_utilities, only: coord_e2f
#endif
        
        character(*), parameter :: rout_name = 'calc_ang_grid_eq_B'
        
        ! input / output
        type(grid_type), intent(inout) :: grid_eq
        type(eq_1_type), intent(in), target :: eq
        logical, intent(in), optional :: only_half_grid
        
        ! local variables
        character(len=max_str_ln) :: err_msg                                    ! error message
        real(dp), allocatable :: r_e_loc(:)                                     ! flux in Equilibrium coords.
        real(dp), pointer :: flux_f(:) => null()                                ! flux that the F uses as normal coord.
        real(dp), pointer :: flux_e(:) => null()                                ! flux that the E uses as normal coord.
        real(dp) :: r_f_factor, r_e_factor                                      ! mult. factors for r_F and r_E
        integer :: pmone                                                        ! plus or minus one
        integer :: id, jd, kd                                                   ! counters
        integer :: n_par_x_loc                                                  ! local n_par_X
        logical :: only_half_grid_loc                                           ! local only_half_grid
        real(dp), allocatable :: theta_f_loc(:,:,:)                             ! local theta_F
        real(dp), allocatable :: zeta_f_loc(:,:,:)                              ! local zeta_F
        
        ! initialize ierr
        ierr = 0
        
        call lvl_ud(1)
        
        ! set local only_half_grid
        only_half_grid_loc = .false.
        if (present(only_half_grid)) only_half_grid_loc = only_half_grid
        
        ! set local n_par_X_rich
        ierr = calc_n_par_x_rich(n_par_x_loc,only_half_grid_loc)
        chckerr('')
        
        ! user output
        call writo('for '//trim(i2str(grid_eq%n(3)))//&
            &' values on normal range '//trim(r2strt(min_r_sol))//'..'//&
            &trim(r2strt(max_r_sol)))
        call writo('for '//trim(i2str(n_par_x))//' values on parallel &
            &range '//trim(r2strt(min_par_x))//'pi..'//&
            &trim(r2strt(max_par_x))//'pi')
        call lvl_ud(1)
        if (only_half_grid_loc) call writo('for this Richardson level, only &
            &the even points are setup up')
        call print_info_eq(n_par_x_loc)
        call lvl_ud(-1)
        
        ! set up flux_E and plus minus one
        ! Note: this routine  is similar to calc_loc_r, but  that routine cannot
        ! be used here because it is possible that the FD quantities are not yet
        ! defined.
        ! choose which equilibrium style is being used:
        !   1:  VMEC
        !   2:  HELENA
        select case (eq_style)
            case (1)                                                            ! VMEC
                if (use_pol_flux_e) then                                        ! normalized poloidal flux
                    flux_e => eq%flux_p_E(:,0)
                else                                                            ! normalized toroidal flux
                    flux_e => eq%flux_t_E(:,0)
                end if
                r_e_factor = max_flux_e
                pmone = -1                                                      ! conversion VMEC LH -> RH coord. system
            case (2)                                                            ! HELENA
                flux_e => eq%flux_p_E(:,0)
                r_e_factor = 2*pi
                pmone = 1
        end select
        
        ! set up flux_F
        if (use_pol_flux_f) then                                                ! poloidal flux / 2pi
            flux_f => eq%flux_p_E(:,0)
            r_f_factor = 2*pi
        else                                                                    ! toroidal flux / 2pi
            flux_f => eq%flux_t_E(:,0)
            r_f_factor = pmone*2*pi                                             ! possible conversion VMEC LH -> RH coord. system
        end if
        
        ! set up parallel angle in  Flux coordinates on equidistant grid        
        ! and use this to calculate the other angle as well
        allocate(theta_f_loc(n_par_x,n_alpha,grid_eq%loc_n_r))
        allocate(zeta_f_loc(n_par_x,n_alpha,grid_eq%loc_n_r))
        if (use_pol_flux_f) then                                                ! parallel angle theta
            ierr = calc_eqd_grid(theta_f_loc,min_par_x*pi,&
                &max_par_x*pi,1,excl_last=.false.)                              ! first index corresponds to parallel angle
            chckerr('')
            do kd = 1,grid_eq%loc_n_r
                zeta_f_loc(:,:,kd) = pmone*eq%q_saf_E(kd,0)*&
                    &theta_f_loc(:,:,kd)
            end do
            do jd = 1,n_alpha
                zeta_f_loc(:,jd,:) = zeta_f_loc(:,jd,:) + alpha(jd)
            end do
        else                                                                    ! parallel angle zeta
            ierr = calc_eqd_grid(zeta_f_loc,min_par_x*pi,&
                &max_par_x*pi,1,excl_last=.false.)                              ! first index corresponds to parallel angle
            chckerr('')
            do kd = 1,grid_eq%loc_n_r
                theta_f_loc(:,:,kd) = pmone*eq%rot_t_E(kd,0)*&
                    &zeta_f_loc(:,:,kd)
            end do
            do jd = 1,n_alpha
                theta_f_loc(:,jd,:) = theta_f_loc(:,jd,:) - alpha(jd)
            end do
        end if
        
        ! set up grid_eq angular F coordinates
        if (only_half_grid_loc) then
            do id = eq_jobs_lims(1,eq_job_nr),eq_jobs_lims(2,eq_job_nr)
                grid_eq%theta_F(id-eq_jobs_lims(1,eq_job_nr)+1,:,:) = &
                    &theta_f_loc(2*id,:,:)
                grid_eq%zeta_F(id-eq_jobs_lims(1,eq_job_nr)+1,:,:) = &
                    &zeta_f_loc(2*id,:,:)
            end do
        else
            grid_eq%theta_F = theta_f_loc(eq_jobs_lims(1,eq_job_nr):&
                &eq_jobs_lims(2,eq_job_nr),:,:)
            grid_eq%zeta_F = zeta_f_loc(eq_jobs_lims(1,eq_job_nr):&
                &eq_jobs_lims(2,eq_job_nr),:,:)
        end if
        
        ! allocate local r_E
        allocate(r_e_loc(size(flux_f)))
        
        ! convert  Flux  coordinates  to  Equilibrium  coordinates  (use
        ! custom flux_E and  flux_F because the Flux  quantities are not
        ! yet calculated)
        call writo('convert F to E coordinates')
        call lvl_ud(1)
        ierr = coord_f2e(grid_eq,grid_eq%loc_r_F,grid_eq%theta_F,&
            &grid_eq%zeta_F,r_e_loc,grid_eq%theta_E,grid_eq%zeta_E,&
            &r_f_array=flux_f/r_f_factor,r_e_array=flux_e/r_e_factor)
        chckerr('')
        call lvl_ud(-1)
        
        ! test whether r_E_loc indeed corresponds to loc_r_E of eq grid
        if (maxval(abs(grid_eq%loc_r_E-r_e_loc)).gt.10*tol_zero) then
            ierr = 1
            err_msg = 'loc_r_E of equilibrium grid is not recovered'
            chckerr(err_msg)
        end if
        
#if ldebug
        if (debug_calc_ang_grid_eq_b) then
            ! test whether F variables recovered
            deallocate(theta_f_loc,zeta_f_loc)
            allocate(theta_f_loc(grid_eq%n(1),grid_eq%n(2),grid_eq%loc_n_r))
            allocate(zeta_f_loc(grid_eq%n(1),grid_eq%n(2),grid_eq%loc_n_r))
            ierr = coord_e2f(grid_eq,grid_eq%loc_r_E,grid_eq%theta_E,&
                &grid_eq%zeta_E,grid_eq%loc_r_F,theta_f_loc,zeta_f_loc,&
                &r_e_array=flux_e/r_e_factor,r_f_array=flux_f/r_f_factor)
            chckerr('')
            
            ! plot difference
            call plot_diff_hdf5(grid_eq%theta_F,theta_f_loc,'TEST_theta_F',&
                &grid_eq%n,[0,0,grid_eq%i_min-1],'test whether F variable are &
                &recovered',output_message=.true.)
            call plot_diff_hdf5(grid_eq%zeta_F,zeta_f_loc,'TEST_zeta_F',&
                &grid_eq%n,[0,0,grid_eq%i_min-1],'test whether F variable are &
                &recovered',output_message=.true.)
        end if
#endif
        
        ! deallocate local variables
        deallocate(r_e_loc)
        nullify(flux_f,flux_e)
        
        call lvl_ud(-1)
    end function calc_ang_grid_eq_b
    
    integer function redistribute_output_grid(grid,grid_out,no_outer_trim) &
        &result(ierr)
        use grid_utilities, only: find_compr_range
        use mpi_utilities, only: redistribute_var, get_ser_var
        use grid_vars, only: n_r_sol
        use num_vars, only: n_procs
        
        character(*), parameter :: rout_name = 'redistribute_output_grid'
        
        ! input / output
        type(grid_type), intent(in) :: grid
        type(grid_type), intent(inout) :: grid_out
        logical, intent(in), optional :: no_outer_trim
        
        ! local variables
        integer :: id                                                           ! counter
        integer :: eq_limits(2)                                                 ! normal limits for equilibrium variables
        integer :: sol_limits(2)                                                ! normal limits for perturbation variables
        integer :: n_out(3)                                                     ! n of grid_out
        integer :: i_lim_tot(2)                                                 ! total limits of grid
        integer :: i_lim_out(2)                                                 ! limits of grid_out
        integer :: lims(2), lims_dis(2)                                         ! limits and distributed limits, taking into account the angular extent
        real(dp), allocatable :: r_f_sol(:)                                     ! perturbation r_F
        real(dp), allocatable :: temp_var(:)                                    ! temporary variable
        integer, allocatable :: temp_lim(:)                                     ! temporary limit
        integer, allocatable :: eq_limits_tot(:,:)                              ! total equilibrium limits
        logical :: no_outer_trim_loc                                            ! local no_outer_trim
        
        ! initialize ierr
        ierr = 0
        
        ! calculate normal range for solution and save in perturbation variables
        allocate(r_f_sol(n_r_sol))
        ierr = calc_norm_range('PB3D_sol',sol_limits=sol_limits,r_f_sol=r_f_sol)
        chckerr('')
        
        ! determine eq_limits: smallest eq range comprising X range
        call find_compr_range(grid%r_F,r_f_sol(sol_limits),eq_limits)
        
        ! get lowest equilibrium limits to be  able to setup an output grid that
        ! starts at index 1
        allocate(eq_limits_tot(2,n_procs))
        do id = 1,2
            ierr = get_ser_var(eq_limits(id:id),temp_lim,scatter=.true.)
            chckerr('')
            eq_limits_tot(id,:) = temp_lim
        end do
        
        ! set up redistributed grid
        no_outer_trim_loc = .false.
        if (present(no_outer_trim)) no_outer_trim_loc = no_outer_trim
        if (no_outer_trim_loc) then
            n_out = grid%n
            i_lim_tot = [1,n_out(3)]
            i_lim_out = eq_limits
        else
            n_out(1:2) = grid%n(1:2)
            n_out(3) = eq_limits_tot(2,n_procs)-eq_limits_tot(1,1)+1
            i_lim_tot = [eq_limits_tot(1,1),eq_limits_tot(2,n_procs)]
            i_lim_out = eq_limits-eq_limits_tot(1,1)+1
        end if
        ierr = grid_out%init(n_out,i_lim_out)
        chckerr('')
        
        ! set up limits taking into account angular extent and temporary var
        lims(1) = product(grid%n(1:2))*(grid%i_min-1)+1
        lims(2) = product(grid%n(1:2))*grid%i_max
        lims_dis(1) = product(grid%n(1:2))*(eq_limits(1)-1)+1
        lims_dis(2) = product(grid%n(1:2))*eq_limits(2)
        allocate(temp_var(lims_dis(2)-lims_dis(1)+1))
        
        ! copy total variables
        ! r_F
        grid_out%r_F = grid%r_F(i_lim_tot(1):i_lim_tot(2))
        ! r_E
        grid_out%r_E = grid%r_E(i_lim_tot(1):i_lim_tot(2))
        
        ! distribute local variables
        ! theta_F
        ierr = redistribute_var(reshape(grid%theta_F,[size(grid%theta_F)]),&
            &temp_var,lims,lims_dis)
        chckerr('')
        grid_out%theta_F = reshape(temp_var,shape(grid_out%theta_F))
        ! zeta_F
        ierr = redistribute_var(reshape(grid%zeta_F,[size(grid%zeta_F)]),&
            &temp_var,lims,lims_dis)
        chckerr('')
        grid_out%zeta_F = reshape(temp_var,shape(grid_out%zeta_F))
        ! theta_E
        ierr = redistribute_var(reshape(grid%theta_E,[size(grid%theta_E)]),&
            &temp_var,lims,lims_dis)
        chckerr('')
        grid_out%theta_E = reshape(temp_var,shape(grid_out%theta_E))
        ! zeta_E
        ierr = redistribute_var(reshape(grid%zeta_E,[size(grid%zeta_E)]),&
            &temp_var,lims,lims_dis)
        chckerr('')
        grid_out%zeta_E = reshape(temp_var,shape(grid_out%zeta_E))
        ! loc_r_F and loc_R_E
        if (grid_out%divided) then
            grid_out%loc_r_F = grid_out%r_F(grid_out%i_min:grid_out%i_max)
            grid_out%loc_r_E = grid_out%r_E(grid_out%i_min:grid_out%i_max)
        end if
    end function redistribute_output_grid
    
    integer function magn_grid_plot(grid) result(ierr)
        use num_vars, only: rank, no_plots, n_theta_plot, n_zeta_plot, &
            &eq_style, min_theta_plot, max_theta_plot, min_zeta_plot, &
            &max_zeta_plot
        use grid_utilities, only: trim_grid, extend_grid_f, calc_xyz_grid
        use vmec_utilities, only: calc_trigon_factors
        
        character(*), parameter :: rout_name = 'magn_grid_plot'
        
        ! input / output
        type(grid_type), intent(in) :: grid
        
        ! local variables
        real(dp), allocatable :: x_1(:,:,:), y_1(:,:,:), z_1(:,:,:)             ! X, Y and Z of surface in Axisymmetric coordinates
        real(dp), allocatable :: x_2(:,:,:), y_2(:,:,:), z_2(:,:,:)             ! X, Y and Z of magnetic field lines in Axisymmetric coordinates
        real(dp), pointer :: x_1_tot(:,:,:) => null()                           ! total X
        real(dp), pointer :: y_1_tot(:,:,:) => null()                           ! total Y
        real(dp), pointer :: z_1_tot(:,:,:) => null()                           ! total Z
        real(dp), pointer :: x_2_tot(:,:,:) => null()                           ! total X
        real(dp), pointer :: y_2_tot(:,:,:) => null()                           ! total Y
        real(dp), pointer :: z_2_tot(:,:,:) => null()                           ! total Z
        type(grid_type) :: grid_ext                                             ! angularly extended grid
        type(grid_type) :: grid_plot                                            ! grid for plotting
        integer :: n_theta_plot_old                                             ! backup of n_theta_plot
        integer :: n_zeta_plot_old                                              ! backup of n_zeta_plot
        real(dp) :: min_theta_plot_old, max_theta_plot_old                      ! backup of min and max_theta_plot
        real(dp) :: min_zeta_plot_old, max_zeta_plot_old                        ! backup of min and max_zeta_plot
        character(len=max_str_ln) :: anim_name                                  ! name of animation
        
        ! initialize ierr
        ierr = 0
        
        ! bypass plots if no_plots
        if (no_plots) return
        
        call writo('Plotting magnetic field and flux surfaces')
        
        call lvl_ud(1)
        
        ! 0. set up variables
        
        ! save n_theta_plot and n_zeta_plot and change them
        n_theta_plot_old = n_theta_plot
        n_zeta_plot_old = n_zeta_plot
        min_theta_plot_old = min_theta_plot
        max_theta_plot_old = max_theta_plot
        min_zeta_plot_old = min_zeta_plot
        max_zeta_plot_old = max_zeta_plot
        n_theta_plot = 40
        n_zeta_plot = 160
        min_theta_plot = 1
        max_theta_plot = 3
        min_zeta_plot = 0
        max_zeta_plot = 2
        
        ! extend grid
        ierr = extend_grid_f(grid,grid_ext,grid_eq=grid)
        chckerr('')
        
        ! restore n_theta_plot and n_zeta_plot
        n_theta_plot = n_theta_plot_old
        n_zeta_plot = n_zeta_plot_old
        min_theta_plot = min_theta_plot_old
        max_theta_plot = max_theta_plot_old
        min_zeta_plot = min_zeta_plot_old
        max_zeta_plot = max_zeta_plot_old
        
        ! trim extended grid into plot grid
        ierr = trim_grid(grid_ext,grid_plot)
        chckerr('')
        
        ! if VMEC, calculate trigonometric factors of plot grid
        if (eq_style.eq.1) then
            ierr = calc_trigon_factors(grid_plot%theta_E,grid_plot%zeta_E,&
                &grid_plot%trigon_factors)
            chckerr('')
        end if
        
        ! set animation name
        anim_name = 'Magnetic field in flux surfaces'
        
        ! 1. plot flux surfaces
        call writo('writing flux surfaces')
        
        ! calculate X_1,Y_1 and Z_1
        allocate(x_1(grid_plot%n(1),grid_plot%n(2),grid_plot%loc_n_r))
        allocate(y_1(grid_plot%n(1),grid_plot%n(2),grid_plot%loc_n_r))
        allocate(z_1(grid_plot%n(1),grid_plot%n(2),grid_plot%loc_n_r))
        ierr = calc_xyz_grid(grid,grid_plot,x_1,y_1,z_1)
        chckerr('')
        
        ! dealloc grid
        call grid_plot%dealloc()
        
        ! get pointers to full X, Y and Z
        ! The reason for this is that the plot  is not as simple as usual, so no
        ! divided  plots  are used,  and  also  efficiency  is not  the  biggest
        ! priority. Therefore,  all the plotting  of the  local is handled  by a
        ! single process, the master.
        call get_full_xyz(x_1,y_1,z_1,x_1_tot,y_1_tot,z_1_tot,'flux surfaces')
        
        ! 2. plot field lines
        call writo('writing field lines')
        
        ! trim grid into plot grid
        ierr = trim_grid(grid,grid_plot)
        chckerr('')
        
        ! if VMEC, calculate trigonometric factors of plot grid
        if (eq_style.eq.1) then
            ierr = calc_trigon_factors(grid_plot%theta_E,grid_plot%zeta_E,&
                &grid_plot%trigon_factors)
            chckerr('')
        end if
        
        ! calculate X_2,Y_2 and Z_2
        allocate(x_2(grid_plot%n(1),grid_plot%n(2),grid_plot%loc_n_r))
        allocate(y_2(grid_plot%n(1),grid_plot%n(2),grid_plot%loc_n_r))
        allocate(z_2(grid_plot%n(1),grid_plot%n(2),grid_plot%loc_n_r))
        ierr = calc_xyz_grid(grid,grid_plot,x_2,y_2,z_2)
        chckerr('')
        
        ! dealloc grids
        call grid_plot%dealloc()
        call grid_ext%dealloc()
        
        ! get pointers to full X, Y and Z
        ! The reason for this is that the plot  is not as simple as usual, so no
        ! divided  plots  are used,  and  also  efficiency  is not  the  biggest
        ! priority. Therefore,  all the plotting  of the  local is handled  by a
        ! single process, the master.
        call get_full_xyz(x_2,y_2,z_2,x_2_tot,y_2_tot,z_2_tot,'field lines')
        
        ierr = magn_grid_plot_hdf5(x_1_tot,x_2_tot,y_1_tot,y_2_tot,&
            &z_1_tot,z_2_tot,anim_name)
        chckerr('')
        
        ! clean up
        deallocate(x_1,y_1,z_1)
        deallocate(x_2,y_2,z_2)
        nullify(x_1_tot,y_1_tot,z_1_tot)
        nullify(x_2_tot,y_2_tot,z_2_tot)
        
        call lvl_ud(-1)
        
        call writo('Done plotting magnetic field and flux surfaces')
    contains
        ! get pointer to full plotting variables X, Y and Z
        subroutine get_full_xyz(X,Y,Z,X_tot,Y_tot,Z_tot,merge_name)
            use mpi_utilities, only: get_ser_var
            
            ! input / output
            real(dp), intent(in), target :: x(:,:,:), y(:,:,:), z(:,:,:)        ! X, Y and Z of either flux surfaces or magnetic field lines
            real(dp), intent(inout), pointer :: x_tot(:,:,:)                    ! pointer to full X
            real(dp), intent(inout), pointer :: y_tot(:,:,:)                    ! pointer to full Y
            real(dp), intent(inout), pointer :: z_tot(:,:,:)                    ! pointer to full Z
            character(len=*) :: merge_name                                      ! name of variable to be merged
            
            ! local variables
            real(dp), allocatable :: ser_xyz_loc(:)                             ! serial copy of XYZ_loc
            integer, allocatable :: tot_dim(:)                                  ! total dimensions for plot 
            
            ! merge plots for flux surfaces if more than one process
            if (grid%divided) then                                              ! merge local plots
                ! user output
                call writo('Merging local plots for '//merge_name)
                
                ! get total dimension
                ierr = get_ser_var([size(x,3)],tot_dim)
                chckerr('')
                
                ! allocate total X, Y and Z (should have same sizes)
                if (rank.eq.0) then
                    allocate(x_tot(size(x,1),size(x,2),sum(tot_dim)))
                    allocate(y_tot(size(x,1),size(x,2),sum(tot_dim)))
                    allocate(z_tot(size(x,1),size(x,2),sum(tot_dim)))
                end if
                
                allocate(ser_xyz_loc(size(x(:,:,1))*sum(tot_dim)))
                ierr = get_ser_var(reshape(x,[size(x)]),ser_xyz_loc)
                chckerr('')
                if (rank.eq.0) x_tot = reshape(ser_xyz_loc,shape(x_tot))
                ierr = get_ser_var(reshape(y,[size(y)]),ser_xyz_loc)
                chckerr('')
                if (rank.eq.0) y_tot = reshape(ser_xyz_loc,shape(y_tot))
                ierr = get_ser_var(reshape(z,[size(z)]),ser_xyz_loc)
                chckerr('')
                if (rank.eq.0) z_tot = reshape(ser_xyz_loc,shape(z_tot))
            else                                                                ! just point
                x_tot => x
                y_tot => y
                z_tot => z
            end if
        end subroutine get_full_xyz
        
        ! Plot with HDF5
        integer function magn_grid_plot_hdf5(X_1,X_2,Y_1,Y_2,Z_1,Z_2,&
            &anim_name) result(ierr)
            use hdf5_ops, only: open_hdf5_file, add_hdf5_item, &
                &print_hdf5_top, print_hdf5_geom, print_hdf5_3d_data_item, &
                &print_hdf5_grid, close_hdf5_file
            use hdf5_vars, only: dealloc_xml_str, &
                & xml_str_type, hdf5_file_type
            use rich_vars, only: rich_lvl
            use grid_vars, only: n_alpha
            
            character(*), parameter :: rout_name = 'magn_grid_plot_HDF5'
          
            ! input / output
            real(dp), intent(in), pointer :: x_1(:,:,:), y_1(:,:,:), z_1(:,:,:) ! X, Y and Z of surface in Axisymmetric coordinates
            real(dp), intent(in), pointer :: x_2(:,:,:), y_2(:,:,:), z_2(:,:,:) ! X, Y and Z of magnetic field lines in Axisymmetric coordinates
            character(len=*), intent(in) :: anim_name                           ! name of animation
            
            ! local variables
            character(len=max_str_ln) :: plot_title(2)                          ! plot title for flux surface and for field lines
            character(len=max_str_ln) :: file_name                              ! name of file
            integer :: id, jd                                                   ! counter
            type(hdf5_file_type) :: file_info                                   ! file info
            type(xml_str_type), allocatable :: grids(:)                         ! grid in spatial collection (flux surface, field line)
            type(xml_str_type), allocatable :: space_col_grids(:)               ! grid with space collection at different times
            type(xml_str_type) :: time_col_grid                                 ! grid with time collection
            type(xml_str_type) :: top                                           ! topology
            type(xml_str_type) :: xyz(3)                                        ! data items for geometry
            type(xml_str_type) :: geom                                          ! geometry
            integer :: loc_dim(3,2)                                             ! local dimensions (flux surfaces, field lines)
            integer :: n_r                                                      ! nr. of normal points
            
            ! initialize ierr
            ierr = 0
            
            ! only master
            if (rank.eq.0) then
                ! user output
                call writo('Drawing animation with HDF5')
                call lvl_ud(1)
                
                ! set up loc_dim and n_r
                loc_dim(:,1) = [size(x_1,1),size(x_1,2),1]
                loc_dim(:,2) = [size(x_2,1),1,1]
                n_r = size(x_1,3) - 1                                           ! should be same for all other X_i, Y_i and Z_i
                
                ! set up plot titles and file name
                plot_title(1) = 'Magnetic Flux Surfaces'
                plot_title(2) = 'Magnetic Field Lines'
                file_name = 'field_lines_in_flux_surfaces_R_'//&
                    &trim(i2str(rich_lvl))
                
                ! open HDF5 file
                ierr = open_hdf5_file(file_info,trim(file_name),&
                    &descr=anim_name,ind_plot=.true.)
                chckerr('')
                
                ! create grid  for flux surface and all field  lines and one for
                ! time collection
                allocate(grids(1+n_alpha))
                allocate(space_col_grids(n_r))
                
                ! loop over all normal points
                do id = 1,n_r
                    ! A. start with flux surface
                    
                    ! print topology
                    call print_hdf5_top(top,2,loc_dim(:,1))
                    
                    ! print data item for X
                    ierr = print_hdf5_3d_data_item(xyz(1),file_info,&
                        &'X_surf_'//trim(i2str(id)),x_1(:,:,id:id),loc_dim(:,1))
                    chckerr('')
                    
                    ! print data item for Y
                    ierr = print_hdf5_3d_data_item(xyz(2),file_info,&
                        &'Y_surf_'//trim(i2str(id)),y_1(:,:,id:id),loc_dim(:,1))
                    chckerr('')
                    
                    ! print data item for Z
                    ierr = print_hdf5_3d_data_item(xyz(3),file_info,&
                        &'Z_surf_'//trim(i2str(id)),z_1(:,:,id:id),loc_dim(:,1))
                    chckerr('')
                    
                    ! print geometry with X, Y and Z data item
                    call print_hdf5_geom(geom,2,xyz,.true.)
                    
                    ! create a grid with the topology and the geometry
                    ierr = print_hdf5_grid(grids(1),plot_title(1),1,&
                        &grid_top=top,grid_geom=geom,reset=.true.)
                    chckerr('')
                    
                    ! B. end with magnetic field
                    
                    do jd = 1,n_alpha                                           ! iterate over all field lines
                        ! print topology
                        call print_hdf5_top(top,2,loc_dim(:,2))
                        
                        ! print data item for X of this field line
                        ierr = print_hdf5_3d_data_item(xyz(1),file_info,&
                            &'X_field_'//trim(i2str(id))//'_'//trim(i2str(jd)),&
                            &x_2(:,jd:jd,id+1:id+1),loc_dim(:,2))
                        chckerr('')
                        
                        ! print data item for Y of this field line
                        ierr = print_hdf5_3d_data_item(xyz(2),file_info,&
                            &'Y_field_'//trim(i2str(id))//'_'//trim(i2str(jd)),&
                            &y_2(:,jd:jd,id+1:id+1),loc_dim(:,2))
                        chckerr('')
                        
                        ! print data item for Z of this field line
                        ierr = print_hdf5_3d_data_item(xyz(3),file_info,&
                            &'Z_field_'//trim(i2str(id))//'_'//trim(i2str(jd)),&
                            &z_2(:,jd:jd,id+1:id+1),loc_dim(:,2))
                        chckerr('')
                        
                        ! print geometry with X, Y and Z data item
                        call print_hdf5_geom(geom,2,xyz,.true.)
                        
                        ! create a grid with the topology and the geometry
                        ierr = print_hdf5_grid(grids(jd+1),plot_title(2),1,&
                            &grid_top=top,grid_geom=geom,reset=.true.)
                        chckerr('')
                    end do
                    
                    ! C.  merge flux  surface and magnetic  field in  to spatial
                    ! collection
                    ierr = print_hdf5_grid(space_col_grids(id),&
                        &'spatial collection',3,grid_time=id*1._dp,&
                        &grid_grids=grids,reset=.true.)
                    chckerr('')
                end do
                
                ! create grid collection from individual grids and reset them
                ierr = print_hdf5_grid(time_col_grid,'time collection',2,&
                    &grid_grids=space_col_grids,reset=.true.)
                chckerr('')
                
                ! add collection grid to HDF5 file and reset it
                ierr = add_hdf5_item(file_info,time_col_grid,reset=.true.)
                chckerr('')
                
                ! close HDF5 file
                ierr = close_hdf5_file(file_info)
                chckerr('')
                
                call lvl_ud(-1)
                
                ! clean up
                do id = 1,2
                    call dealloc_xml_str(grids(id))
                end do
                call dealloc_xml_str(space_col_grids)
                call dealloc_xml_str(time_col_grid)
                call dealloc_xml_str(top)
                do id = 1,3
                    call dealloc_xml_str(xyz(id))
                end do
                call dealloc_xml_str(geom)
            end if
        end function magn_grid_plot_hdf5
    end function magn_grid_plot
    
    integer function print_output_grid(grid,grid_name,data_name,rich_lvl,&
        &par_div,remove_previous_arrs) result(ierr)
        use num_vars, only: pb3d_name, eq_jobs_lims, eq_job_nr
        use hdf5_ops, only: print_hdf5_arrs
        use hdf5_vars, only: var_1d_type, &
            &max_dim_var_1d
        use grid_utilities, only: trim_grid
        
        character(*), parameter :: rout_name = 'print_output_grid'
        
        ! input / output
        type(grid_type), intent(in) :: grid
        character(len=*), intent(in) :: grid_name
        character(len=*), intent(in) :: data_name
        integer, intent(in), optional :: rich_lvl
        logical, intent(in), optional :: par_div
        logical, intent(in), optional :: remove_previous_arrs
        
        ! local variables
        integer :: n_tot(3)                                                     ! total n
        integer :: par_id(2)                                                    ! local parallel interval
        type(grid_type) :: grid_trim                                            ! trimmed grid
        type(var_1d_type), allocatable, target :: grid_1d(:)                    ! 1D equivalent of grid variables
        type(var_1d_type), pointer :: grid_1d_loc => null()                     ! local element in grid_1D
        integer :: id                                                           ! counter
        logical :: par_div_loc                                                  ! local par_div
        
        ! initialize ierr
        ierr = 0
        
        ! user output
        call writo('Write '//trim(grid_name)//' grid variables to output &
            &file')
        call lvl_ud(1)
        
        ! trim grid
        ierr = trim_grid(grid,grid_trim)
        chckerr('')
        
        ! set local par_div
        par_div_loc = .false.
        if (present(par_div)) par_div_loc = par_div
        
        ! set total n and parallel interval
        n_tot = grid_trim%n
        par_id = [1,n_tot(1)]
        if (grid_trim%n(1).gt.0 .and. par_div_loc) then                         ! total grid includes all equilibrium jobs
            n_tot(1) = maxval(eq_jobs_lims)-minval(eq_jobs_lims)+1
            par_id = eq_jobs_lims(:,eq_job_nr)
        end if
        
        ! Set up the 1D equivalents of the perturbation variables
        allocate(grid_1d(max_dim_var_1d))
        
        ! set up variables grid_1D
        id = 1
        
        ! n
        grid_1d_loc => grid_1d(id); id = id+1
        grid_1d_loc%var_name = 'n'
        allocate(grid_1d_loc%tot_i_min(1),grid_1d_loc%tot_i_max(1))
        allocate(grid_1d_loc%loc_i_min(1),grid_1d_loc%loc_i_max(1))
        grid_1d_loc%tot_i_min = [1]
        grid_1d_loc%tot_i_max = [3]
        grid_1d_loc%loc_i_min = [1]
        grid_1d_loc%loc_i_max = [3]
        allocate(grid_1d_loc%p(3))
        grid_1d_loc%p = n_tot
        
        ! r_F
        grid_1d_loc => grid_1d(id); id = id+1
        grid_1d_loc%var_name = 'r_F'
        allocate(grid_1d_loc%tot_i_min(1),grid_1d_loc%tot_i_max(1))
        allocate(grid_1d_loc%loc_i_min(1),grid_1d_loc%loc_i_max(1))
        grid_1d_loc%tot_i_min = [1]
        grid_1d_loc%tot_i_max = [n_tot(3)]
        grid_1d_loc%loc_i_min = [grid_trim%i_min]
        grid_1d_loc%loc_i_max = [grid_trim%i_max]
        allocate(grid_1d_loc%p(size(grid_trim%loc_r_F)))
        grid_1d_loc%p = grid_trim%loc_r_F
        
        ! r_E
        grid_1d_loc => grid_1d(id); id = id+1
        grid_1d_loc%var_name = 'r_E'
        allocate(grid_1d_loc%tot_i_min(1),grid_1d_loc%tot_i_max(1))
        allocate(grid_1d_loc%loc_i_min(1),grid_1d_loc%loc_i_max(1))
        grid_1d_loc%tot_i_min = [1]
        grid_1d_loc%tot_i_max = [n_tot(3)]
        grid_1d_loc%loc_i_min = [grid_trim%i_min]
        grid_1d_loc%loc_i_max = [grid_trim%i_max]
        allocate(grid_1d_loc%p(size(grid_trim%loc_r_E)))
        grid_1d_loc%p = grid_trim%loc_r_E
        
        ! only for 3D grids
        if (product(grid%n(1:2)).ne.0) then
            ! theta_F
            grid_1d_loc => grid_1d(id); id = id+1
            grid_1d_loc%var_name = 'theta_F'
            allocate(grid_1d_loc%tot_i_min(3),grid_1d_loc%tot_i_max(3))
            allocate(grid_1d_loc%loc_i_min(3),grid_1d_loc%loc_i_max(3))
            grid_1d_loc%tot_i_min = [1,1,1]
            grid_1d_loc%tot_i_max = n_tot
            grid_1d_loc%loc_i_min = [par_id(1),1,grid_trim%i_min]
            grid_1d_loc%loc_i_max = [par_id(2),n_tot(2),grid_trim%i_max]
            allocate(grid_1d_loc%p(size(grid_trim%theta_F)))
            grid_1d_loc%p = reshape(grid_trim%theta_F,[size(grid_trim%theta_F)])
            
            ! theta_E
            grid_1d_loc => grid_1d(id); id = id+1
            grid_1d_loc%var_name = 'theta_E'
            allocate(grid_1d_loc%tot_i_min(3),grid_1d_loc%tot_i_max(3))
            allocate(grid_1d_loc%loc_i_min(3),grid_1d_loc%loc_i_max(3))
            grid_1d_loc%tot_i_min = [1,1,1]
            grid_1d_loc%tot_i_max = n_tot
            grid_1d_loc%loc_i_min = [par_id(1),1,grid_trim%i_min]
            grid_1d_loc%loc_i_max = [par_id(2),n_tot(2),grid_trim%i_max]
            allocate(grid_1d_loc%p(size(grid_trim%theta_E)))
            grid_1d_loc%p = reshape(grid_trim%theta_E,[size(grid_trim%theta_E)])
            
            ! zeta_F
            grid_1d_loc => grid_1d(id); id = id+1
            grid_1d_loc%var_name = 'zeta_F'
            allocate(grid_1d_loc%tot_i_min(3),grid_1d_loc%tot_i_max(3))
            allocate(grid_1d_loc%loc_i_min(3),grid_1d_loc%loc_i_max(3))
            grid_1d_loc%tot_i_min = [1,1,1]
            grid_1d_loc%tot_i_max = n_tot
            grid_1d_loc%loc_i_min = [par_id(1),1,grid_trim%i_min]
            grid_1d_loc%loc_i_max = [par_id(2),n_tot(2),grid_trim%i_max]
            allocate(grid_1d_loc%p(size(grid_trim%zeta_F)))
            grid_1d_loc%p = reshape(grid_trim%zeta_F,[size(grid_trim%zeta_F)])
            
            ! zeta_E
            grid_1d_loc => grid_1d(id); id = id+1
            grid_1d_loc%var_name = 'zeta_E'
            allocate(grid_1d_loc%tot_i_min(3),grid_1d_loc%tot_i_max(3))
            allocate(grid_1d_loc%loc_i_min(3),grid_1d_loc%loc_i_max(3))
            grid_1d_loc%tot_i_min = [1,1,1]
            grid_1d_loc%tot_i_max = n_tot
            grid_1d_loc%loc_i_min = [par_id(1),1,grid_trim%i_min]
            grid_1d_loc%loc_i_max = [par_id(2),n_tot(2),grid_trim%i_max]
            allocate(grid_1d_loc%p(size(grid_trim%zeta_E)))
            grid_1d_loc%p = reshape(grid_trim%zeta_E,[size(grid_trim%zeta_E)])
        end if
        
        ! write
        ierr = print_hdf5_arrs(grid_1d(1:id-1),pb3d_name,&
            &'grid_'//trim(data_name),rich_lvl=rich_lvl,&
            &ind_print=.not.grid_trim%divided,&
            &remove_previous_arrs=remove_previous_arrs)
        chckerr('')
        
        ! clean up
        call grid_trim%dealloc()
        
        ! clean up
        nullify(grid_1d_loc)
        
        ! user output
        call lvl_ud(-1)
    end function print_output_grid
end module grid_ops