vac_vars.f90

Go to the documentation of this file.

!------------------------------------------------------------------------------!
!------------------------------------------------------------------------------!
module vac_vars
#include <PB3D_macros.h>
    use strumpackdensepackage
    use str_utilities
    use messages
    use num_vars, only: dp, max_name_ln, iu, weight_dp, max_str_ln
    use grid_vars, only: grid_type
    
    implicit none
    
    private
    public copy_vac, set_loc_lims, in_context
#if ldebug
    public n_alloc_vacs
#endif
    
    ! global variables
    integer, parameter :: bs = 16
#if ldebug
    integer :: n_alloc_vacs
#endif
    
    type, public :: vac_type
        integer :: style
        integer :: prim_x
        integer :: ctxt_hg
        integer :: n_bnd
        integer :: bs
        integer :: mpi_comm
        integer :: desc_h(blacsctxtsize)
        integer :: desc_g(blacsctxtsize)
        integer :: n_p(2)
        integer :: n_ang(2)
        integer :: ind_p(2)
        integer :: n_loc(2)
        integer, allocatable :: sec_x(:)
        integer, allocatable :: lims_c(:,:)
        integer, allocatable :: lims_r(:,:)
        real(dp) :: jq
        real(dp), allocatable :: ang(:,:)
        real(dp), allocatable :: norm(:,:)
        real(dp), allocatable :: dnorm(:,:)
        real(dp), allocatable :: h_fac(:,:)
        real(dp), allocatable :: x_vec(:,:)
        real(dp), allocatable :: h(:,:)
        real(dp), allocatable :: g(:,:)
        complex(dp), allocatable :: res(:,:)
#if ldebug
        real(dp) :: estim_mem_usage
#endif
    contains
        procedure :: init => init_vac
        procedure :: dealloc => dealloc_vac
    end type
    
contains
    integer function init_vac(vac,style,n_bnd,prim_X,n_ang,jq) result(ierr)
        use mpi
        use num_vars, only: n_procs, rank
        use x_vars, only: n_mod_x
#if ldebug
        use num_vars, only: print_mem_usage
#endif
        
        character(*), parameter :: rout_name = 'init_vac'
        
        ! input / output
        class(vac_type), intent(inout) :: vac
        integer, intent(in) :: style
        integer, intent(in) :: n_bnd
        integer, intent(in) :: prim_x
        integer, intent(in) :: n_ang(2)
        real(dp), intent(in) :: jq
        
        ! local variables
        character(len=max_str_ln) :: err_msg                                    ! error message
        integer :: n_dims                                                       ! number of dimensions to be saved in x_vec and norm
        integer :: jd                                                           ! counter
        integer, allocatable :: proc_map(:,:)                                   ! process map
        
        ! initialize ierr
        ierr = 0
        
#if ldebug
        ! initialize memory usage
        if (print_mem_usage) vac%estim_mem_usage = 0._dp
#endif
        
        ! test
        if (product(n_ang).ne.n_bnd) then
            ierr = 1
            err_msg = 'The total number of angles along B, '//&
                &trim(i2str(product(n_ang)))//', is not equal to n_bnd = '//&
                &trim(i2str(n_bnd))
            chckerr(err_msg)
        end if
        
        ! set some variables
        vac%bs = bs
        if (n_procs.eq.1) vac%bs = n_bnd
        vac%jq = jq
        vac%n_ang = n_ang
        
        ! Initialize the BLACS grid for H and G
        vac%prim_X = prim_x
        vac%n_bnd = n_bnd
        vac%n_p(1)=floor(sqrt(n_procs*1._dp))
        vac%n_p(2)=n_procs/vac%n_p(1)
        allocate(proc_map(vac%n_p(1),vac%n_p(2)))
        proc_map = transpose(reshape(&
            &[(jd+n_procs-1-product(vac%n_p),jd=1,product(vac%n_p))],vac%n_p))
        call blacs_get(0,0,vac%ctxt_HG)
        call blacs_gridmap(vac%ctxt_HG,proc_map,vac%n_p(1),vac%n_p(1),&
            &vac%n_p(2))                                                        ! last MPI proc is also last process in map
        call blacs_gridinfo(vac%ctxt_HG,vac%n_p(1),vac%n_p(2),vac%ind_p(1),&
            &vac%ind_p(2))                                                      ! get nr. of rows and columns and local index
        if(minval(vac%ind_p).ge.0) then                                         ! only the ranks that participate
            do jd = 1,2                                                         ! loop over rows, then columns
                vac%n_loc(jd) = numroc(vac%n_bnd,vac%bs,vac%ind_p(jd),0,&
                    &vac%n_p(jd))                                               ! get local number of rows, then columns
            end do
            call descinit(vac%desc_H,n_bnd,n_bnd,vac%bs,vac%bs,0,0,vac%ctxt_HG,&
                &max(1,vac%n_loc(1)),ierr)
            chckerr('descinit failed')
            call descinit(vac%desc_G,n_bnd,n_bnd,vac%bs,vac%bs,0,0,vac%ctxt_HG,&
                &max(1,vac%n_loc(1)),ierr)
            chckerr('descinit failed')
            call mpi_comm_split(mpi_comm_world,1,rank,vac%MPI_Comm,ierr)
            chckerr('')
        else
            vac%n_loc = [0,0]
            call mpi_comm_split(mpi_comm_world,mpi_undefined,rank,&
                &vac%MPI_Comm,ierr)
            chckerr('')
        end if
        
        ! set limits
        call set_loc_lims(vac%n_loc(1),vac%bs,vac%ind_p(1),vac%n_p(1),&
            &vac%lims_r)
        call set_loc_lims(vac%n_loc(2),vac%bs,vac%ind_p(2),vac%n_p(2),&
            &vac%lims_c)
        
        ! set n_dims
        vac%style = style
        select case (vac%style)
            case (1)                                                            ! field-line 3-D
                n_dims = 3
            case (2)                                                            ! axisymmetric
                n_dims = 2
            case default
                ierr = 1
                err_msg = 'No vacuum style '//trim(i2str(vac%style))//&
                    &' possible'
                chckerr(err_msg)
        end select
        
        ! allocate normal and position vector
        allocate(vac%norm(vac%n_bnd,n_dims))
        allocate(vac%x_vec(vac%n_bnd,n_dims))
        
        ! allocate variables specific to style
        select case (vac%style)
            case (1)                                                            ! field-line 3-D
                allocate(vac%h_fac(vac%n_bnd,4))
            case (2)                                                            ! axisymmetric
                allocate(vac%ang(n_ang(1),n_ang(2)))
                allocate(vac%dnorm(vac%n_bnd,n_dims))
        end select
        
        ! allocate vacuum response if last process
        if (rank.eq.n_procs-1) allocate(vac%res(n_mod_x,n_mod_x))
        
        ! allocate H and G
        allocate(vac%G(vac%n_loc(1),vac%n_loc(2)))
        allocate(vac%H(vac%n_loc(1),vac%n_loc(2)))
        vac%G = 0._dp
        vac%H = 0._dp
        
#if ldebug
        ! set estimated memory usage
        if (print_mem_usage) vac%estim_mem_usage = &
            &vac%estim_mem_usage + size(vac%G) + n_mod_x**2
        
        ! increment n_alloc_vacs
        n_alloc_vacs = n_alloc_vacs + 1
        
        ! print memory usage
        if (print_mem_usage) call writo('[rank '//trim(i2str(rank))//&
            &' - Expected memory usage of vac: '//&
            &trim(r2strt(vac%estim_mem_usage*weight_dp*2))//' kB]',alert=.true.)
#endif
    end function init_vac
    
    subroutine set_loc_lims(n_loc,bs,ind_p,n_p,lims)
        ! input / output
        integer, intent(in) :: n_loc
        integer, intent(in) :: bs
        integer, intent(in) :: ind_p
        integer, intent(in) :: n_p
        integer, intent(inout), allocatable :: lims(:,:)
        
        ! local variables
        integer :: jd                                                           ! counter
        
        ! set limits
        if (n_loc.gt.0) then
            allocate(lims(2,ceiling(n_loc*1._dp/bs)))
            do jd = 1,size(lims,2)
                lims(1,jd) = indxl2g((jd-1)*bs+1,bs,ind_p,0,n_p)
                lims(2,jd) = indxl2g(min(jd*bs,n_loc),bs,ind_p,0,n_p)
            end do
        else
            ! these processes don't play along
            allocate(lims(2,1))
            lims(:,1) = [0,-1]
        end if
    end subroutine set_loc_lims
    
    subroutine dealloc_vac(vac)
#if ldebug
        use num_vars, only: rank, print_mem_usage
#endif
        
        ! input / output
        class(vac_type), intent(inout) :: vac
        
#if ldebug
        ! local variables
        integer :: mem_diff                                                     ! difference in memory
        real(dp) :: estim_mem_usage                                             ! estimated memory usage
        
        ! memory usage before deallocation
        if (print_mem_usage) then
            mem_diff = get_mem_usage()
            estim_mem_usage = vac%estim_mem_usage
        end if
#endif
        
        ! free grid
        if (in_context(vac%ctxt_HG)) call blacs_gridexit(vac%ctxt_HG)
        
        ! deallocate allocatable variables
        call dealloc_vac_final(vac)
        
#if ldebug
        ! decrement n_alloc_vacs
        n_alloc_vacs = n_alloc_vacs - 1
        
        ! memory usage difference after deallocation
        if (print_mem_usage) then
            mem_diff = mem_diff - get_mem_usage()
            call writo('[Rank '//trim(i2str(rank))//' - liberated '//&
                &trim(i2str(mem_diff))//'kB deallocating vac ('//&
                &trim(i2str(nint(100*mem_diff/&
                &(estim_mem_usage*weight_dp*2))))//&
                &'% of estimated)]',alert=.true.)
        end if
#endif
    contains
        ! Note: intent(out) automatically deallocates the variable
        subroutine dealloc_vac_final(vac)
            ! input / output
            type(vac_type), intent(out) :: vac                                  ! vacuum variables to be deallocated
        end subroutine dealloc_vac_final
    end subroutine dealloc_vac
    
    integer function copy_vac(vac,vac_copy) result(ierr)
        use num_vars, only: rank, n_procs
        
        ! input / output
        class(vac_type), intent(in) :: vac
        class(vac_type), intent(inout) :: vac_copy
        
        character(*), parameter :: rout_name = 'copy_vac'
        
        ! initialize ierr
        ierr = 0
        
        ! reallocate
        ierr = vac_copy%init(vac%style,vac%n_bnd,vac%prim_X,vac%n_ang,&
            &vac%jq)
        chckerr('')
        
        ! copy normal and position vector
        vac_copy%norm = vac%norm
        vac_copy%x_vec = vac%x_vec
        
        ! copy variables specific to style
        select case (vac%style)
            case (1)                                                            ! field-line 3-D
                vac_copy%h_fac = vac%h_fac
            case (2)                                                            ! axisymmetric
                vac_copy%dnorm = vac%dnorm
                vac_copy%ang = vac%ang
        end select
        
        ! copy vacuum response
        if (rank.eq.n_procs-1) vac_copy%res = vac%res
        
        ! copy H and G
        vac_copy%H = vac%H
        vac_copy%G = vac%G
    end function copy_vac
    
    logical function in_context(ctxt) result(res)
        ! input / output
        integer, intent(in) :: ctxt
        
        ! local variables
        integer :: n_p(2)                                                       ! nr. of processes in grid
        integer :: ind_p(2)                                                     ! index of local process in grid
        
        call blacs_gridinfo(ctxt,n_p(1),n_p(2),ind_p(1),ind_p(2)) 
        res = ind_p(1).ge.0
    end function in_context
end module vac_vars