X_vars.f90

Go to the documentation of this file.

!------------------------------------------------------------------------------!
!------------------------------------------------------------------------------!
module x_vars
#include <PB3D_macros.h>
    use str_utilities
    use messages
    use num_vars, only: dp, max_name_ln, iu, weight_dp
    use grid_vars, only: grid_type
    
    implicit none
    
    private
    public set_nm_x, set_nn_mod, &
        &n_mod_x, prim_x, min_sec_x, max_sec_x, min_nm_x, min_n_x, max_n_x, &
        &min_m_x, max_m_x, min_r_sol, max_r_sol, mds_x, mds_sol
#if ldebug
    public n_alloc_x_1s, n_alloc_x_2s
#endif
    
    type, public :: modes_type
        integer, allocatable :: n(:,:)
        integer, allocatable :: m(:,:)
        integer, allocatable :: sec(:,:)
    contains
        procedure :: dealloc => dealloc_mds
    end type
    
    type, public :: x_1_type
        integer :: n_mod
        integer :: lim_sec_x(2)
        integer, allocatable :: n(:,:)
        integer, allocatable :: m(:,:)
        complex(dp), allocatable :: u_0(:,:,:,:)
        complex(dp), allocatable :: u_1(:,:,:,:)
        complex(dp), allocatable :: du_0(:,:,:,:)
        complex(dp), allocatable :: du_1(:,:,:,:)
#if ldebug
        real(dp) :: estim_mem_usage
#endif
    contains
        procedure :: init => init_x_1
        procedure :: copy => copy_x_1
        procedure :: dealloc => dealloc_x_1
    end type
    
    type, public :: x_2_type
        integer :: n_mod(2)
        integer :: lim_sec_x(2,2)
        integer, allocatable :: n_1(:,:)
        integer, allocatable :: n_2(:,:)
        integer, allocatable :: m_1(:,:)
        integer, allocatable :: m_2(:,:)
        complex(dp), allocatable :: pv_0(:,:,:,:)
        complex(dp), allocatable :: pv_1(:,:,:,:)
        complex(dp), allocatable :: pv_2(:,:,:,:)
        complex(dp), allocatable :: kv_0(:,:,:,:)
        complex(dp), allocatable :: kv_1(:,:,:,:)
        complex(dp), allocatable :: kv_2(:,:,:,:)
#if ldebug
        real(dp) :: estim_mem_usage
#endif
    contains
        procedure :: init => init_x_2
        procedure :: copy => copy_x_2
        procedure :: dealloc => dealloc_x_2
    end type
    
    ! interfaces
    
    interface set_nm_x
        module procedure set_nm_x_1
        module procedure set_nm_x_2
    end interface
    
    ! global variables
    type(modes_type) :: mds_x
    type(modes_type) :: mds_sol
    integer :: prim_x
    integer :: min_sec_x
    integer :: max_sec_x
    integer :: n_mod_x
    integer :: min_nm_x = 5
    integer, allocatable :: min_n_x(:)
    integer, allocatable :: max_n_x(:)
    integer, allocatable :: min_m_x(:)
    integer, allocatable :: max_m_x(:)
    real(dp) :: min_r_sol
    real(dp) :: max_r_sol
#if ldebug
    integer :: n_alloc_x_1s
    integer :: n_alloc_x_2s
#endif
    
contains
    subroutine set_nm_x_1(mds,grid_X,lim_sec_X_o,n_X_loc,m_X_loc,lim_sec_X_i)
        
        ! input / output
        type(modes_type), intent(in) :: mds
        type(grid_type), intent(in) :: grid_X
        integer, intent(inout) :: lim_sec_X_o(2)
        integer, intent(inout), allocatable :: n_X_loc(:,:)
        integer, intent(inout), allocatable :: m_X_loc(:,:)
        integer, intent(in), optional :: lim_sec_X_i(2)
        
        ! set output lim_sec_X
        lim_sec_x_o = [1,n_mod_x]
        if (present(lim_sec_x_i)) lim_sec_x_o = lim_sec_x_i
        
        ! set n and m
        allocate(n_x_loc(grid_x%loc_n_r,lim_sec_x_o(2)-lim_sec_x_o(1)+1))
        allocate(m_x_loc(grid_x%loc_n_r,lim_sec_x_o(2)-lim_sec_x_o(1)+1))
        n_x_loc = mds%n(grid_x%i_min:grid_x%i_max,&
            &lim_sec_x_o(1):lim_sec_x_o(2))
        m_x_loc = mds%m(grid_x%i_min:grid_x%i_max,&
            &lim_sec_x_o(1):lim_sec_x_o(2))
    end subroutine set_nm_x_1
    subroutine set_nm_x_2(mds,grid_X,lim_sec_X_o,n_X_1,m_X_1,n_X_2,m_X_2,&
        &lim_sec_X_i)
        
        ! input / output
        type(modes_type), intent(in) :: mds
        type(grid_type), intent(in) :: grid_X
        integer, intent(inout) :: lim_sec_X_o(2,2)
        integer, intent(inout), allocatable :: n_X_1(:,:)
        integer, intent(inout), allocatable :: m_X_1(:,:)
        integer, intent(inout), allocatable :: n_X_2(:,:)
        integer, intent(inout), allocatable :: m_X_2(:,:)
        integer, intent(in), optional :: lim_sec_X_i(2,2)
        
        ! call vectorial version
        if (present(lim_sec_x_i)) then
            call set_nm_x(mds,grid_x,lim_sec_x_o(:,1),n_x_1,m_x_1,&
                &lim_sec_x_i(:,1))
            call set_nm_x(mds,grid_x,lim_sec_x_o(:,2),n_x_2,m_x_2,&
                &lim_sec_x_i(:,2))
        else
            call set_nm_x(mds,grid_x,lim_sec_x_o(:,1),n_x_1,m_x_1)
            call set_nm_x(mds,grid_x,lim_sec_x_o(:,2),n_x_2,m_x_2)
        end if
    end subroutine set_nm_x_2
    
    subroutine init_x_1(X,mds,grid_X,lim_sec_X)
#if ldebug
        use num_vars, only: print_mem_usage, rank
#endif
        
        ! input / output
        class(x_1_type), intent(inout) :: X
        type(modes_type), intent(in) :: mds
        type(grid_type), intent(in) :: grid_X
        integer, intent(in), optional :: lim_sec_X(2)
        
        ! local variables
        integer :: loc_n_r                                                      ! local nr. of normal points
        integer :: n_par, n_geo                                                 ! tot. nr. of angular points in parallel and geodesic direction
        
        ! set local variables
        loc_n_r = grid_x%loc_n_r
        n_par = grid_x%n(1)
        n_geo = grid_x%n(2)
        
#if ldebug
        ! initialize memory usage
        if (print_mem_usage) x%estim_mem_usage = 0._dp
#endif
        
        ! set mode numbers
        call set_nm_x(mds,grid_x,x%lim_sec_X,x%n,x%m,lim_sec_x)
        
        ! set n_mod
        x%n_mod = size(x%n,2)
        
        ! allocate U_0
        allocate(x%U_0(n_par,n_geo,loc_n_r,x%n_mod))
        
        ! allocate U_1
        allocate(x%U_1(n_par,n_geo,loc_n_r,x%n_mod))
        
        ! allocate DU_0
        allocate(x%DU_0(n_par,n_geo,loc_n_r,x%n_mod))
        
        ! allocate DU_1
        allocate(x%DU_1(n_par,n_geo,loc_n_r,x%n_mod))
        
#if ldebug
        ! set estimated memory usage
        if (print_mem_usage) x%estim_mem_usage = x%estim_mem_usage + &
            &(n_par*n_geo*loc_n_r)*(x%n_mod*4)
        
        ! increment n_alloc_X_1s
        n_alloc_x_1s = n_alloc_x_1s + 1
        
        ! print memory usage
        if (print_mem_usage) call writo('[rank '//trim(i2str(rank))//&
            &' - Expected memory usage of X_1: '//&
            &trim(r2strt(x%estim_mem_usage*weight_dp*2))//' kB]',alert=.true.)
#endif
    end subroutine init_x_1
    subroutine init_x_2(X,mds,grid_X,lim_sec_X,is_field_averaged)
#if ldebug
        use num_vars, only: print_mem_usage, rank
#endif
        
        ! input / output
        class(x_2_type), intent(inout) :: X
        type(modes_type), intent(in) :: mds
        type(grid_type), intent(in) :: grid_X
        integer, intent(in), optional :: lim_sec_X(2,2)
        logical, intent(in), optional :: is_field_averaged
        
        ! local variables
        integer :: loc_n_r                                                      ! local nr. of normal points
        integer :: n_par, n_geo                                                 ! tot. nr. of angular points in parallel and geodesic direction
        integer :: nn_mod                                                       ! either n_mod^2, n_mod*(n_mod+1)/2 or 0
        
        ! set local variables
        n_par = grid_x%n(1)
        if (present(is_field_averaged)) then
            if (is_field_averaged) n_par = 1
        end if
        n_geo = grid_x%n(2)
        loc_n_r = grid_x%loc_n_r
        
#if ldebug
        ! initialize memory usage
        if (print_mem_usage) x%estim_mem_usage = 0._dp
#endif
        
        ! set mode numbers
        call set_nm_x(mds,grid_x,x%lim_sec_X,x%n_1,x%m_1,x%n_2,x%m_2,lim_sec_x)
        
        ! set n_mod
        x%n_mod(1) = size(x%n_1,2)
        x%n_mod(2) = size(x%n_2,2)
        
        ! set nnmod for symmetric quantities
        nn_mod = set_nn_mod(.true.,lim_sec_x)
        
        ! allocate PV_i
        allocate(x%PV_0(n_par,n_geo,loc_n_r,nn_mod))                            ! symmetric
        allocate(x%PV_1(n_par,n_geo,loc_n_r,product(x%n_mod)))                  ! not symmetric
        allocate(x%PV_2(n_par,n_geo,loc_n_r,nn_mod))                            ! symmetric
        
        ! allocate KV_i
        allocate(x%KV_0(n_par,n_geo,loc_n_r,nn_mod))                            ! symmetric
        allocate(x%KV_1(n_par,n_geo,loc_n_r,product(x%n_mod)))                  ! not symmetric
        allocate(x%KV_2(n_par,n_geo,loc_n_r,nn_mod))                            ! symmetric
        
#if ldebug
        ! set estimated memory usage
        if (print_mem_usage) x%estim_mem_usage = &
            &x%estim_mem_usage + (n_par*n_geo*loc_n_r)*&
            &(nn_mod*4+product(x%n_mod)*2) + product(x%n_mod)
        
        ! increment n_alloc_X_2s
        n_alloc_x_2s = n_alloc_x_2s + 1
        
        ! print memory usage
        if (print_mem_usage) call writo('[rank '//trim(i2str(rank))//&
            &' - Expected memory usage of X_2: '//&
            &trim(r2strt(x%estim_mem_usage*weight_dp*2))//' kB]',alert=.true.)
#endif
    end subroutine init_x_2
    
    subroutine copy_x_1(X_i,mds,grid_i,X_o)
        use grid_vars, only: grid_type
        
        ! input / output
        class(x_1_type), intent(in) :: X_i
        type(modes_type), intent(in) :: mds
        type(grid_type), intent(in) :: grid_i
        type(x_1_type), intent(inout) :: X_o
        
        call x_o%init(mds,grid_i,lim_sec_x=x_i%lim_sec_X)
        
        ! U_i
        x_o%U_0 = x_i%U_0
        x_o%U_1 = x_i%U_1
        
        ! DU_i
        x_o%DU_0 = x_i%DU_0
        x_o%DU_1 = x_i%DU_1
    end subroutine copy_x_1
    
    subroutine copy_x_2(X_i,mds,grid_i,X_o)
        use grid_vars, only: grid_type
        
        ! input / output
        class(x_2_type), intent(in) :: X_i
        type(modes_type), intent(in) :: mds
        type(grid_type), intent(in) :: grid_i
        type(x_2_type), intent(inout) :: X_o
        
        call x_o%init(mds,grid_i,lim_sec_x=x_i%lim_sec_X,&
            &is_field_averaged=size(x_i%PV_0,1).eq.1)
        
        ! PV_i
        x_o%PV_0 = x_i%PV_0
        x_o%PV_1 = x_i%PV_1
        x_o%PV_2 = x_i%PV_2
        
        ! KV_i
        x_o%KV_0 = x_i%KV_0
        x_o%KV_1 = x_i%KV_1
        x_o%KV_2 = x_i%KV_2
    end subroutine copy_x_2
    
    integer function set_nn_mod(sym,lim_sec_X) result(nn_mod)
        ! input / output
        logical, intent(in) :: sym
        integer, intent(in), optional :: lim_sec_x(2,2)
        
        ! local variables
        integer :: id, jd                                                       ! counters
        integer :: lim_sec_x_loc(2,2)                                           ! local version of lim_sec_X
        
        ! set local lim_sec_X
        lim_sec_x_loc(:,1) = [1,n_mod_x]
        lim_sec_x_loc(:,2) = [1,n_mod_x]
        if (present(lim_sec_x)) lim_sec_x_loc = lim_sec_x
        
        if (sym) then
            ! set nnmod for symmetric quantities: discard indices above diagonal
            nn_mod = 0
            do jd = lim_sec_x_loc(1,2),lim_sec_x_loc(2,2)
                do id = lim_sec_x_loc(1,1),lim_sec_x_loc(2,1)
                    if (id.ge.jd) nn_mod = nn_mod + 1
                end do
            end do
        else
            ! set nn_mod for asymmetric quantities: don't discard anything
            nn_mod = product(lim_sec_x_loc(2,:)-lim_sec_x_loc(1,:)+1)
        end if
    end function set_nn_mod
    
    subroutine dealloc_mds(mds)
        ! input / output
        class(modes_type), intent(inout) :: mds
        
        ! deallocate allocatable variables
        call dealloc_mds_final(mds)
    contains
        ! Note: intent(out) automatically deallocates the variable
        subroutine dealloc_mds_final(mds)
            ! input / output
            type(modes_type), intent(out) :: mds                                ! modes variables to be deallocated
        end subroutine dealloc_mds_final
    end subroutine dealloc_mds
    
    subroutine dealloc_x_1(X)
#if ldebug
        use num_vars, only: rank, print_mem_usage
#endif
        
        ! input / output
        class(x_1_type), intent(inout) :: X
        
#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 = x%estim_mem_usage
        end if
#endif
        
        ! deallocate allocatable variables
        call dealloc_x_1_final(x)
        
#if ldebug
        ! decrement n_alloc_X_1s
        n_alloc_x_1s = n_alloc_x_1s - 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 X_1 ('//&
                &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_x_1_final(X)
            ! input / output
            type(x_1_type), intent(out) :: X                                    ! equilibrium to be deallocated
        end subroutine dealloc_x_1_final
    end subroutine dealloc_x_1
    subroutine dealloc_x_2(X)
#if ldebug
        use num_vars, only: rank, print_mem_usage
#endif
        
        ! input / output
        class(x_2_type), intent(inout) :: X
        
#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 = x%estim_mem_usage
        end if
#endif
        
        ! deallocate allocatable variables
        call dealloc_x_2_final(x)
        
#if ldebug
        ! decrement n_alloc_X_2s
        n_alloc_x_2s = n_alloc_x_2s - 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 X_2 ('//&
                &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_x_2_final(X)
            ! input / output
            type(x_2_type), intent(out) :: X                                    ! equilibrium to be deallocated
        end subroutine dealloc_x_2_final
    end subroutine dealloc_x_2
end module x_vars