driver_sol.f90

Go to the documentation of this file.

!------------------------------------------------------------------------------!
!------------------------------------------------------------------------------!
module driver_sol
#include <PB3D_macros.h>
#include <wrappers.h>
    use str_utilities
    use output_ops
    use messages
    use num_vars, only: dp, pi, max_str_ln, iu
    use grid_vars, only: grid_type
    use x_vars, only: x_2_type, modes_type, &
        &mds_x, mds_sol
    use vac_vars, only: vac_type
    use sol_vars, only: sol_type
    
    implicit none
    private
    public run_driver_sol
    
    ! global variables
#if ldebug
    logical :: debug_run_driver_sol = .false.                                   
    logical :: debug_interp_V = .false.                                         
    integer :: ivs_stats(3)
#endif
    
contains
    integer function run_driver_sol(grid_eq,grid_X,grid_sol,X,vac,sol) &
        &result(ierr)
        
        use num_vars, only: ev_style, eq_style, rich_restart_lvl, rank, &
            &n_procs, x_grid_style, jump_to_sol
        use grid_vars, only: n_r_sol, n_alpha
        use pb3d_ops, only: reconstruct_pb3d_grid, reconstruct_pb3d_sol
        use slepc_ops, only: solve_ev_system_slepc
        use grid_ops, only: calc_norm_range, setup_grid_sol, &
            &print_output_grid, redistribute_output_grid
        use sol_ops, only: print_output_sol
        use rich_vars, only: rich_lvl
        use rich_ops, only: calc_rich_ex
        use vac_ops, only: calc_vac_res, print_output_vac
        use grid_utilities, only: trim_grid
        use x_ops, only: setup_modes, redistribute_output_x
#if ldebug
        use x_ops, only: print_debug_x_2
#endif
        
        character(*), parameter :: rout_name = 'run_driver_sol'
        
        ! input / output
        type(grid_type), intent(in) :: grid_eq
        type(grid_type), intent(in), target :: grid_x
        type(grid_type), intent(inout) :: grid_sol
        type(x_2_type), intent(in) :: x
        type(vac_type), intent(inout) :: vac
        type(sol_type), intent(inout) :: sol
        
        ! local variables
        type(grid_type) :: grid_sol_trim                                        ! trimmed solution grid
        type(grid_type) :: grid_x_rdst                                          ! redistributed perturbation grid
        type(grid_type) :: grid_x_sol                                           ! perturbation grid with solution normal part
        type(x_2_type) :: x_rdst                                                ! redistributed X
        type(x_2_type) :: x_sol                                                 ! interpolated X
        integer :: sol_limits(2)                                                ! min. and max. index of sol grid for this process
        integer :: rich_lvl_name                                                ! either the Richardson level or zero, to append to names
        real(dp), allocatable :: r_f_sol(:)                                     ! normal points in solution grid
        logical :: do_vac_ops                                                   ! whether specific calculations for vacuum are necessary
        character(len=max_str_ln) :: err_msg                                    ! error message
        
        ! initialize ierr
        ierr = 0
        
        ! set up solution grid if first level
        if (rich_lvl.eq.rich_restart_lvl) then
            ! Divide solution grid under group processes, calculating the limits
            ! and the normal coordinate.
            allocate(r_f_sol(n_r_sol))
            ierr = calc_norm_range('PB3D_sol',sol_limits=sol_limits,&
                &r_f_sol=r_f_sol)
            chckerr('')
            
            if (rich_lvl.eq.1) then
                call writo('Set up solution grid')
                call lvl_ud(1)
                
                call writo('Calculate the grid')
                call lvl_ud(1)
                ierr = setup_grid_sol(grid_eq,grid_x,grid_sol,r_f_sol,&
                    &sol_limits)
                chckerr('')
                call lvl_ud(-1)
                
                call writo('Write to output file')
                call lvl_ud(1)
                ierr = print_output_grid(grid_sol,'solution','sol',&
                    &remove_previous_arrs=(jump_to_sol.and.&
                    &(x_grid_style.eq.1 .or. x_grid_style.eq.3)))
                chckerr('')
                call lvl_ud(-1)
                
                ! set up modes
                ierr = setup_modes(mds_sol,grid_eq,grid_sol,plot_name='sol')
                chckerr('')
                
                call lvl_ud(-1)
            else
                ! restore solution grid and trim it
                ierr = reconstruct_pb3d_grid(grid_sol,'sol',&
                    &grid_limits=sol_limits)
                chckerr('')
                ierr = trim_grid(grid_sol,grid_sol_trim)
                chckerr('')
                
                ! set up modes
                ierr = setup_modes(mds_sol,grid_eq,grid_sol,plot_name='sol')
                chckerr('')
                
                ! reconstruct solution on trimmed grid
                ierr = reconstruct_pb3d_sol(mds_sol,grid_sol_trim,sol,'sol',&
                    &rich_lvl=rich_lvl-1)
                chckerr('')
                
                ! clean up
                call grid_sol_trim%dealloc()
            end if
            
            deallocate(r_f_sol)
        end if
        
        ! set up  whether Richardson level  has to be  appended to the  name and
        ! whether to do vacuum operations
        select case (eq_style) 
            case (1)                                                            ! VMEC
                rich_lvl_name = rich_lvl                                        ! append richardson level
                do_vac_ops = .true.
            case (2)                                                            ! HELENA
                rich_lvl_name = 0                                               ! do not append name
                if (rich_lvl.eq.rich_restart_lvl) then
                    do_vac_ops = .true.
                else
                    do_vac_ops = .false.
                end if
        end select
        
        if (do_vac_ops) then
            ! calculate vacuum
            ierr = calc_vac_res(mds_sol,vac)
            chckerr('')
            
            call writo('Write to output file')
            call lvl_ud(1)
            chckerr('')
            if (rank.eq.n_procs-1) then
                ierr = print_output_vac(vac,'vac',rich_lvl=rich_lvl_name)
                chckerr('')
            end if
            call lvl_ud(-1)
        end if
        
        ! initialize interpolated X
        ierr = grid_x_sol%init([1,n_alpha,grid_sol%n(3)],&
            &[grid_sol%i_min,grid_sol%i_max],grid_sol%divided)
        chckerr('')
        grid_x_sol%r_F = grid_sol%r_F
        grid_x_sol%r_E = grid_sol%r_E
        grid_x_sol%loc_r_F = grid_sol%loc_r_F
        grid_x_sol%loc_r_E = grid_sol%loc_r_E
        
        
        select case (x_grid_style)
            case (1,3)                                                          ! equilibrium or enriched
                ! user output
                call writo('Redistribute and interpolate the perturbation &
                    &variables to solution grid')
                call lvl_ud(1)
                
                ! redistribute grid and X variables
                ierr = redistribute_output_grid(grid_x,grid_x_rdst)
                chckerr('')
                ierr = redistribute_output_x(mds_x,grid_x,grid_x_rdst,x,x_rdst)
                chckerr('')
                
                ! initialize
                call x_sol%init(mds_sol,grid_x_sol,is_field_averaged=.true.)
                
                ! interpolate
                ierr = interp_v(mds_x,grid_x_rdst,x_rdst,mds_sol,grid_x_sol,&
                    &x_sol)
                chckerr('')
                
                ! clean up
                call x_rdst%dealloc()
                call grid_x_rdst%dealloc()
                
                call lvl_ud(-1)
            case (2)                                                            ! solution
                ! user output
                call writo('Copy the perturbation variables to solution grid')
                call lvl_ud(1)
                
                ! copy
                call x%copy(mds_sol,grid_x,x_sol)
                
                call lvl_ud(-1)
        end select
        
#if ldebug
        ! write integrated  field-aligned tensorial  perturbation quantities
        ! to output and plot
        if (debug_run_driver_sol) then
            ierr = print_debug_x_2(mds_sol,grid_x_sol,x_sol)
            chckerr('')
        end if
#endif
        
        ! clean up
        call grid_x_sol%dealloc()
        
        ! solve the system
        call writo('Solving the system')
        call lvl_ud(1)
        select case (ev_style)
            case(1)                                                             ! SLEPC solver for EV problem
                ! solve the system
                ierr = solve_ev_system_slepc(mds_sol,grid_sol,x_sol,vac,sol)
                chckerr('')
            case default
                err_msg = 'No EV solver style associated with '//&
                    &trim(i2str(ev_style))
                ierr = 1
                chckerr(err_msg)
        end select
        call lvl_ud(-1)
        call writo('System solved')
        
        ! write solution variables to output
        ierr = print_output_sol(grid_sol,sol,'sol',rich_lvl=rich_lvl,&
            &remove_previous_arrs=&
            &(jump_to_sol.and.(x_grid_style.eq.1 .or. x_grid_style.eq.3)))
        chckerr('')
        
        ! calculate Richardson extrapolation factors if necessary
        call calc_rich_ex(sol%val)
        
        ! clean up
        call x_sol%dealloc()
    end function run_driver_sol
    
    integer function interp_v(mds_i,grid_i,X_i,mds_o,grid_o,X_o) result(ierr)
        use x_utilities, only: is_necessary_x, trim_modes
        use x_vars, only: n_mod_x
        use num_vars, only: rank
        use eq_vars, only: max_flux_f
        use num_utilities, only: c
        use sol_utilities, only: interp_v_spline
!#if ldebug
        !use num_vars, only: ex_plot_style
!#endif
        
        character(*), parameter :: rout_name = 'interp_V'
        
        ! input / output
        type(modes_type), intent(in), target :: mds_i
        type(grid_type), intent(in) :: grid_i
        type(x_2_type), intent(in), target :: x_i
        type(modes_type), intent(in), target :: mds_o
        type(grid_type), intent(in) :: grid_o
        type(x_2_type), intent(inout), target :: x_o
        
        ! local variables
        integer :: kd                                                           ! counter
        integer :: k, m                                                         ! counters for mode numbers
        integer :: c_loc_i(2)                                                   ! local input c for symmetric and asymmetric variables
        integer :: c_loc_o(2)                                                   ! local output c for symmetric and asymmetric variables
        integer :: n_mod_tot                                                    ! total amount of modes
        integer :: kdl_i(2), kdl_o(2)                                           ! limits on normal index for a mode combination
        integer :: id_lim_i(2), id_lim_o(2)                                     ! limits on total modes
        integer :: ivs_stat(6)                                                  ! local stats for interp_V_spline
        integer, pointer :: sec_i_loc(:,:), sec_o_loc(:,:)                      ! pointers to secondary mode variables
        real(dp), pointer :: r_i_loc(:), r_o_loc(:)                             ! local r_i and r_o for a mode combination
        complex(dp), pointer :: v_i(:,:,:), v_o(:,:,:)                          ! pointers to input and output PV_i and KV_i
        logical :: calc_this(2)                                                 ! whether this combination needs to be calculated
        logical :: extrap = .true.                                              ! whether extrapolation is used
        character(len=max_str_ln) :: err_msg                                    ! error message
#if ldebug
        integer :: km_id
        integer, allocatable :: norm_ext_i(:,:)                                 ! normal extent for input quantity mode combinations
        real(dp), allocatable :: r_loc_tot(:,:,:)                               ! r_i_loc and r_o_loc for all combinations
        character(len=max_str_ln) :: ivs_stats_names(3)                         ! names used in interp_V_spline statistics
        !complex(dp), allocatable :: V_plot(:,:)                                 ! for debug plotting of interpolated V
#endif
        
        ! initialize ierr
        ierr = 0
        
        ! test
        if (grid_i%n(2).ne.grid_o%n(2)) then
            ierr = 1
            err_msg = 'input and output grid are not compatible in the &
                &geodesic coordinate: '//trim(i2str(grid_i%n(2)))//' vs. '//&
                &trim(i2str(grid_o%n(2)))
            chckerr(err_msg)
        end if
        
        ! initialize to zero
        x_o%PV_0 = 0._dp
        x_o%PV_1 = 0._dp
        x_o%PV_2 = 0._dp
        x_o%KV_0 = 0._dp
        x_o%KV_1 = 0._dp
        x_o%KV_2 = 0._dp
        
        ! limit input modes to output modes
        ! (X grid should comprise sol grid)
        ierr = trim_modes(mds_i,mds_o,id_lim_i,id_lim_o)
        chckerr('')
        sec_i_loc => mds_i%sec(id_lim_i(1):id_lim_i(2),:)
        sec_o_loc => mds_o%sec(id_lim_o(1):id_lim_o(2),:)
        
        ! select all mode number combinations of interpolated grid
        n_mod_tot = size(sec_o_loc,1)
#if ldebug
        allocate(r_loc_tot(2,n_mod_tot**2,3))
        allocate(norm_ext_i(n_mod_tot,n_mod_tot))
        km_id = 0
        norm_ext_i = 0
#endif
        do m = 1,n_mod_tot
            do k = 1,n_mod_tot
#if ldebug
                ! test whether input and output mode numbers are consistent
                if (sec_i_loc(k,1).ne.sec_o_loc(k,1) .or. &
                    &sec_i_loc(m,1).ne.sec_o_loc(m,1)) then
                    ierr = 1
                    err_msg = 'For ('//trim(i2str(k))//','//trim(i2str(m))//&
                        &'), no consistency for modes in grid_i and grid_o'
                    chckerr(err_msg)
                end if
#endif
                
                ! set input and output grid normal limits for mode pair (k,m)
                kdl_i(1) = max(sec_i_loc(k,2),sec_i_loc(m,2))
                kdl_i(2) = min(sec_i_loc(k,3),sec_i_loc(m,3))
                kdl_o(1) = max(sec_o_loc(k,2),sec_o_loc(m,2))
                kdl_o(2) = min(sec_o_loc(k,3),sec_o_loc(m,3))
                
                ! limit to grid ranges
                kdl_i(1) = max(kdl_i(1),grid_i%i_min)
                kdl_i(2) = min(kdl_i(2),grid_i%i_max)
                kdl_o(1) = max(kdl_o(1),grid_o%i_min)
                kdl_o(2) = min(kdl_o(2),grid_o%i_max)
                
                ! limit output range to input range if not extrapolating
                if (.not.extrap) then
                    do kd = kdl_o(1),kdl_o(2)
                        if (grid_o%r_F(kd).ge.grid_i%r_F(kdl_i(1))) exit
                    end do
                    kdl_o(1) = kd
                    
                    do kd = kdl_o(2),kdl_o(1),-1
                        if (grid_o%r_F(kd).le.grid_i%r_F(kdl_i(2))) exit
                    end do
                    kdl_o(2) = kd
                end if
                
                ! cycle if mode combination does not exist anywhere
                if (kdl_i(1).gt.kdl_i(2) .or. kdl_o(1).gt.kdl_o(2)) cycle
                
                ! convert limits to local
                kdl_i = kdl_i - grid_i%i_min + 1
                kdl_o = kdl_o - grid_o%i_min + 1
                
                ! get input ranges for this mode number
                r_i_loc => grid_i%loc_r_F(kdl_i(1):kdl_i(2))
                r_o_loc => grid_o%loc_r_F(kdl_o(1):kdl_o(2))
                
                ! check whether mode combination needs to be calculated
                calc_this(1) = is_necessary_x(.true.,&
                    &[sec_o_loc(k,1),sec_o_loc(m,1)])
                calc_this(2) = is_necessary_x(.false.,&
                    &[sec_o_loc(k,1),sec_o_loc(m,1)])
                
#if ldebug
                km_id = km_id + 1
                r_loc_tot(:,km_id,1) = [r_i_loc(1),r_i_loc(size(r_i_loc))]
                r_loc_tot(:,km_id,2) = [r_o_loc(1),r_o_loc(size(r_o_loc))]
                r_loc_tot(:,km_id,3) = km_id
                norm_ext_i(k,m) = size(r_i_loc)
                if (debug_interp_v) write(*,*) 'k, m', k, m, 'of', n_mod_tot, &
                        &'with calc_this = ', calc_this
#endif
                
                ! set up c_loc for input and output
                c_loc_i(1) = c([sec_i_loc(k,4),sec_i_loc(m,4)],.true.,n_mod_x)
                c_loc_i(2) = c([sec_i_loc(k,4),sec_i_loc(m,4)],.false.,n_mod_x)
                c_loc_o(1) = c([sec_o_loc(k,4),sec_o_loc(m,4)],.true.,n_mod_x)
                c_loc_o(2) = c([sec_o_loc(k,4),sec_o_loc(m,4)],.false.,n_mod_x)
                
                if (calc_this(1)) then
                    v_i => x_i%PV_0(:,:,kdl_i(1):kdl_i(2),c_loc_i(1))
                    v_o => x_o%PV_0(:,:,kdl_o(1):kdl_o(2),c_loc_o(1))
                    ierr = interp_v_spline(v_i,v_o,r_i_loc,r_o_loc,extrap,&
                        &ivs_stat(1))
                    chckerr('')
                    
                    v_i => x_i%PV_2(:,:,kdl_i(1):kdl_i(2),c_loc_i(1))
                    v_o => x_o%PV_2(:,:,kdl_o(1):kdl_o(2),c_loc_o(1))
                    ierr = interp_v_spline(v_i,v_o,r_i_loc,r_o_loc,extrap,&
                        &ivs_stat(2))
                    chckerr('')
                    
                    v_i => x_i%KV_0(:,:,kdl_i(1):kdl_i(2),c_loc_i(1))
                    v_o => x_o%KV_0(:,:,kdl_o(1):kdl_o(2),c_loc_o(1))
                    ierr = interp_v_spline(v_i,v_o,r_i_loc,r_o_loc,extrap,&
                        &ivs_stat(3))
                    chckerr('')
                    
                    v_i => x_i%KV_2(:,:,kdl_i(1):kdl_i(2),c_loc_i(1))
                    v_o => x_o%KV_2(:,:,kdl_o(1):kdl_o(2),c_loc_o(1))
                    ierr = interp_v_spline(v_i,v_o,r_i_loc,r_o_loc,extrap,&
                        &ivs_stat(4))
                    chckerr('')
                    
#if ldebug
                    if (debug_interp_v) then
                        !call writo('For [k,m] = ['//&
                            !&trim(i2str(sec_o_loc(k,1)))//','//&
                            !&trim(i2str(sec_o_loc(m,1)))//']:')
                        !call lvl_ud(1)
                        !call writo('input normal range: '//&
                            !&trim(i2str(kdl_i(1)))//'..'//trim(i2str(kdl_i(2))))
                        !call writo('output normal range: '//&
                            !&trim(i2str(kdl_o(1)))//'..'//trim(i2str(kdl_o(2))))
                        !allocate(V_plot(max(size(V_i,3),size(V_o,3)),4))
                        !V_plot(1:size(V_i,3),1) = V_i(1,1,:)
                        !V_plot(1:size(V_o,3),2) = V_o(1,1,:)
                        !V_plot(1:size(V_i,3),3) = r_i_loc/max_flux_F*2*pi
                        !V_plot(1:size(V_o,3),4) = r_o_loc/max_flux_F*2*pi
                        !V_plot(size(V_i,3)+1:size(V_plot,1),1) = &
                            !&V_plot(size(V_i,3),1)
                        !V_plot(size(V_o,3)+1:size(V_plot,1),2) = &
                            !&V_plot(size(V_o,3),2)
                        !V_plot(size(V_i,3)+1:size(V_plot,1),3) = &
                            !&V_plot(size(V_i,3),3)
                        !V_plot(size(V_o,3)+1:size(V_plot,1),4) = &
                            !&V_plot(size(V_o,3),4)
                        !call print_ex_2D(['V_i','V_o'],'RE_V',&
                            !&rp(V_plot(:,1:2)),x=rp(V_plot(:,3:4)))
                        !!call print_ex_2D(['V_i','V_o'],'IM_V',&
                            !!&ip(V_plot(:,1:2)),x=rp(V_plot(:,3:4)))
                        !deallocate(V_plot)
                        !if (ex_plot_style.eq.2) read(*,*)
                        !call lvl_ud(-1)
                    end if
#endif
                end if
                
                if (calc_this(2)) then
                    v_i => x_i%PV_1(:,:,kdl_i(1):kdl_i(2),c_loc_i(2))
                    v_o => x_o%PV_1(:,:,kdl_o(1):kdl_o(2),c_loc_o(2))
                    ierr = interp_v_spline(v_i,v_o,r_i_loc,r_o_loc,extrap,&
                        &ivs_stat(5))
                    chckerr('')
                    
                    v_i => x_i%KV_1(:,:,kdl_i(1):kdl_i(2),c_loc_i(2))
                    v_o => x_o%KV_1(:,:,kdl_o(1):kdl_o(2),c_loc_o(2))
                    ierr = interp_v_spline(v_i,v_o,r_i_loc,r_o_loc,extrap,&
                        &ivs_stat(6))
                    chckerr('')
                end if
#if ldebug
                do kd = 1,size(ivs_stat)
                    !if (ivs_stat(kd).eq.1) call writo('Used direct copy for '//&
                        !&trim(i2str(kd))//' of mode combination ('//&
                        !&trim(i2str(k))//', '//trim(i2str(m))//') at ψ = '//&
                        !&trim(r2strt(r_loc_tot(1,km_id,1)/max_flux_F*2*pi)))
                    ivs_stats(ivs_stat(kd)) = ivs_stats(ivs_stat(kd)) + 1
                end do
#endif
            end do
        end do
        
#if ldebug
        call writo('interp_V_spline statistics:')
        call lvl_ud(1)
        ivs_stats_names(1) = 'direct copy'
        ivs_stats_names(2) = 'linear interpolation'
        ivs_stats_names(3) = 'spline interpolation'
        do kd = 1,3
            call writo(trim(ivs_stats_names(kd))//' was performed '//&
                &trim(i2str(ivs_stats(kd)))//' times: '//&
                &trim(r2strt(100._dp*ivs_stats(kd)/sum(ivs_stats)))//'%')
        end do
        if (ivs_stats(1).gt.0) then
            call writo('It is possible that more secondary modes and/or a &
                &lower max_njq_change are needed.',warning=.true.)
            call writo('If unphysical modes appear, consider changing these &
                &variables.',warning=.true.)
        end if
        call lvl_ud(-1)
        call print_ex_2d([''],'r_i_loc_'//trim(i2str(rank)),&
            &r_loc_tot(:,1:km_id,1)/max_flux_f*2*pi,x=r_loc_tot(:,1:km_id,3),&
            &draw=.false.)
        call draw_ex([''],'r_i_loc_'//trim(i2str(rank)),km_id,1,.false.,&
            &ex_plot_style=1)
        call print_ex_2d([''],'r_o_loc_'//trim(i2str(rank)),&
            &r_loc_tot(:,1:km_id,2)/max_flux_f*2*pi,x=r_loc_tot(:,1:km_id,3),&
            &draw=.false.)
        call draw_ex([''],'r_o_loc_'//trim(i2str(rank)),km_id,1,.false.,&
            &ex_plot_style=1)
        call plot_hdf5('normal extent','norm_ext_i_'//trim(i2str(rank)),1._dp*&
            &reshape(norm_ext_i*1._dp,[n_mod_tot,n_mod_tot,1]))
#endif
        
        ! clean up
        nullify(r_i_loc,r_o_loc)
        nullify(sec_i_loc,sec_o_loc)
        nullify(v_i,v_o)
    end function interp_v
end module driver_sol