#:include "common.fypp"
#:set RANKS = range(1, MAXRANK + 1)
#:set REDRANKS = range(2, MAXRANK + 1)
#:set RC_KINDS_TYPES = REAL_KINDS_TYPES + CMPLX_KINDS_TYPES
submodule (stdlib_stats) stdlib_stats_moment_mask
use, intrinsic:: ieee_arithmetic, only: ieee_value, ieee_quiet_nan
use stdlib_error, only: error_stop
use stdlib_optval, only: optval
implicit none
contains
#:for k1, t1 in RC_KINDS_TYPES
#:for rank in RANKS
#:set RName = rname("moment_mask",rank, t1, k1)
module function ${RName}$(x, order, dim, center, mask) result(res)
${t1}$, intent(in) :: x${ranksuffix(rank)}$
integer, intent(in) :: order
integer, intent(in) :: dim
${t1}$, intent(in), optional :: center${reduced_shape('x', rank, 'dim')}$
logical, intent(in) :: mask${ranksuffix(rank)}$
${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
integer :: i
real(${k1}$) :: n${reduced_shape('x', rank, 'dim')}$
${t1}$, allocatable :: mean_${ranksuffix(rank-1)}$
n = real(count(mask, dim), ${k1}$)
res = 0
select case(dim)
#:for fi in range(1, rank+1)
case(${fi}$)
if (present(center)) then
do i = 1, size(x, ${fi}$)
res = res + merge( (x${select_subarray(rank, [(fi, 'i')])}$ -&
center)**order,&
#:if t1[0] == 'r'
0._${k1}$,&
#:else
cmplx(0,0,kind=${k1}$),&
#:endif
mask${select_subarray(rank, [(fi, 'i')])}$)
end do
else
allocate(mean_, source = mean(x, ${fi}$, mask))
do i = 1, size(x, ${fi}$)
res = res + merge( (x${select_subarray(rank, [(fi, 'i')])}$ - mean_)**order,&
#:if t1[0] == 'r'
0._${k1}$,&
#:else
cmplx(0,0,kind=${k1}$),&
#:endif
mask${select_subarray(rank, [(fi, 'i')])}$)
end do
deallocate(mean_)
end if
#:endfor
case default
call error_stop("ERROR (moment): wrong dimension")
end select
res = res / n
end function ${RName}$
#:endfor
#:endfor
#:for k1, t1 in INT_KINDS_TYPES
#:for rank in RANKS
#:set RName = rname("moment_mask",rank, t1, k1, 'dp')
module function ${RName}$(x, order, dim, center, mask) result(res)
${t1}$, intent(in) :: x${ranksuffix(rank)}$
integer, intent(in) :: order
integer, intent(in) :: dim
real(dp), intent(in), optional :: center${reduced_shape('x', rank, 'dim')}$
logical, intent(in) :: mask${ranksuffix(rank)}$
real(dp) :: res${reduced_shape('x', rank, 'dim')}$
integer :: i
real(dp) :: n${reduced_shape('x', rank, 'dim')}$
real(dp), allocatable :: mean_${ranksuffix(rank-1)}$
n = real(count(mask, dim), dp)
res = 0
select case(dim)
#:for fi in range(1, rank+1)
case(${fi}$)
if (present(center)) then
do i = 1, size(x, ${fi}$)
res = res + merge((real(x${select_subarray(rank, [(fi, 'i')])}$, dp) -&
center)**order,&
0._dp, mask${select_subarray(rank, [(fi, 'i')])}$)
end do
else
allocate(mean_, source = mean(x, ${fi}$, mask))
do i = 1, size(x, ${fi}$)
res = res + merge((real(x${select_subarray(rank, [(fi, 'i')])}$, dp) - mean_)&
**order,&
0._dp, mask${select_subarray(rank, [(fi, 'i')])}$)
end do
deallocate(mean_)
end if
#:endfor
case default
call error_stop("ERROR (moment): wrong dimension")
end select
res = res / n
end function ${RName}$
#:endfor
#:endfor
end submodule
