mask_mean_dev: addition of a mask in mean API

src/stdlib_experimental_stats.fypp

@@ -15,42 +15,92 @@ module stdlib_experimental_stats
 
     #:for k1, t1 in REAL_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_all_${k1}$_${k1}$(x) result(res)
+        module function mean_${rank}$_all_${k1}$_${k1}$(x, mask) result(res)
           ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          logical, intent(in), optional :: mask
           ${t1}$ :: res
         end function mean_${rank}$_all_${k1}$_${k1}$
       #:endfor
     #:endfor
 
     #:for k1, t1 in INT_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_all_${k1}$_dp(x) result(res)
+        module function mean_${rank}$_all_${k1}$_dp(x, mask) result(res)
           ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          logical, intent(in), optional :: mask
           real(dp) :: res
         end function mean_${rank}$_all_${k1}$_dp
       #:endfor
     #:endfor
 
     #:for k1, t1 in REAL_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_${k1}$_${k1}$(x, dim) result(res)
+        module function mean_${rank}$_${k1}$_${k1}$(x, dim, mask) result(res)
           ${t1}$, intent(in) :: x${ranksuffix(rank)}$
           integer, intent(in) :: dim
+          logical, intent(in), optional :: mask
           ${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
         end function mean_${rank}$_${k1}$_${k1}$
       #:endfor
     #:endfor
 
     #:for k1, t1 in INT_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_${k1}$_dp(x, dim) result(res)
+        module function mean_${rank}$_${k1}$_dp(x, dim, mask) result(res)
           ${t1}$, intent(in) :: x${ranksuffix(rank)}$
           integer, intent(in) :: dim
+          logical, intent(in), optional :: mask
           real(dp) :: res${reduced_shape('x', rank, 'dim')}$
         end function mean_${rank}$_${k1}$_dp
       #:endfor
     #:endfor
 
+
+    #:for k1, t1 in REAL_KINDS_TYPES
+      #:for rank in RANKS
+        module function mean_${rank}$_mask_all_${k1}$_${k1}$(x, mask) result(res)
+          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          logical, intent(in) :: mask${ranksuffix(rank)}$
+          ${t1}$ :: res
+        end function mean_${rank}$_mask_all_${k1}$_${k1}$
+      #:endfor
+    #:endfor
+
+
+    #:for k1, t1 in INT_KINDS_TYPES
+      #:for rank in RANKS
+        module function mean_${rank}$_mask_all_${k1}$_dp(x, mask) result(res)
+          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          logical, intent(in) :: mask${ranksuffix(rank)}$
+          real(dp) :: res
+        end function mean_${rank}$_mask_all_${k1}$_dp
+      #:endfor
+    #:endfor
+
+
+    #:for k1, t1 in REAL_KINDS_TYPES
+      #:for rank in RANKS
+        module function mean_${rank}$_mask_${k1}$_${k1}$(x, dim, mask) result(res)
+          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          integer, intent(in) :: dim
+          logical, intent(in) :: mask${ranksuffix(rank)}$
+          ${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
+        end function mean_${rank}$_mask_${k1}$_${k1}$
+      #:endfor
+    #:endfor
+
+
+    #:for k1, t1 in INT_KINDS_TYPES
+      #:for rank in RANKS
+        module function mean_${rank}$_mask_${k1}$_dp(x, dim, mask) result(res)
+          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          integer, intent(in) :: dim
+          logical, intent(in) :: mask${ranksuffix(rank)}$
+          real(dp) :: res${reduced_shape('x', rank, 'dim')}$
+        end function mean_${rank}$_mask_${k1}$_dp
+      #:endfor
+    #:endfor
+
   end interface mean
 
 end module stdlib_experimental_stats

src/stdlib_experimental_stats.md

@@ -8,36 +8,42 @@
 
 ### Description
 
-Returns the mean of all the elements of `array`, or of the elements of `array` along dimension `dim` if provided.
+Returns the mean of all the elements of `array`, or of the elements of `array` along dimension `dim` if provided, and if the corresponding element in `mask` is `true`.
 
 ### Syntax
 
-`result = mean(array)`
+`result = mean(array [, mask])`
 
-`result = mean(array, dim)`
+`result = mean(array, dim [, mask])`
 
 ### Arguments
 
 `array`: Shall be an array of type `integer`, or `real`.
 
 `dim`: Shall be a scalar of type `integer` with a value in the range from 1 to n, where n is the rank of `array`.
 
+`mask` (optional): Shall be of type `logical` and either by a scalar or an array of the same shape as `array`.
+
 ### Return value
 
 If `array` is of type `real`, the result is of the same type as `array`.
 If `array` is of type `integer`, the result is of type `double precision`.
 
 If `dim` is absent, a scalar with the mean of all elements in `array` is returned. Otherwise, an array of rank n-1, where n equals the rank of `array`, and a shape similar to that of `array` with dimension `dim` dropped is returned.
 
+If `mask` is specified, the result is the mean of all elements of `array` corresponding to `true` elements of `mask`. If every element of `mask` is `false`, the result is IEEE `NaN`.
+
 ### Example
 
 ```fortran
 program demo_mean
     use stdlib_experimental_stats, only: mean
     implicit none
     real :: x(1:6) = [ 1., 2., 3., 4., 5., 6. ]
-    print *, mean(x)                            !returns 21.
-    print *, mean( reshape(x, [ 2, 3 ] ))       !returns 21.
-    print *, mean( reshape(x, [ 2, 3 ] ), 1)    !returns [ 3., 7., 11. ]
+    print *, mean(x)                            !returns 3.5
+    print *, mean( reshape(x, [ 2, 3 ] ))       !returns 3.5
+    print *, mean( reshape(x, [ 2, 3 ] ), 1)    !returns [ 1.5, 3.5, 5.5 ]
+    print *, mean( reshape(x, [ 2, 3 ] ), 1,&
+                   reshape(x, [ 2, 3 ] ) > 3.)  !returns [ NaN, 4.0, 5.5 ]
 end program demo_mean
 ```

src/stdlib_experimental_stats_mean.fypp

@@ -5,17 +5,25 @@
 
 submodule (stdlib_experimental_stats) stdlib_experimental_stats_mean
 
+  use, intrinsic:: ieee_arithmetic, only: ieee_value, ieee_quiet_nan
   use stdlib_experimental_error, only: error_stop
+  use stdlib_experimental_optval, only: optval
   implicit none
 
 contains
 
   #:for k1, t1 in REAL_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_all_${k1}$_${k1}$(x) result(res)
+      module function mean_${rank}$_all_${k1}$_${k1}$(x, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        logical, intent(in), optional :: mask
         ${t1}$ :: res
 
+        if (.not.optval(mask, .true.)) then
+          res = ieee_value(res, ieee_quiet_nan)
+          return
+        end if
+
         res = sum(x) / real(size(x, kind = int64), ${k1}$)
 
       end function mean_${rank}$_all_${k1}$_${k1}$
@@ -25,10 +33,16 @@ contains
 
   #:for k1, t1 in INT_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_all_${k1}$_dp(x) result(res)
+      module function mean_${rank}$_all_${k1}$_dp(x, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        logical, intent(in), optional :: mask
         real(dp) :: res
 
+        if (.not.optval(mask, .true.)) then
+          res = ieee_value(res, ieee_quiet_nan)
+          return
+        end if
+
         res = sum(real(x, dp)) / real(size(x, kind = int64), dp)
 
       end function mean_${rank}$_all_${k1}$_dp
@@ -38,11 +52,17 @@ contains
 
   #:for k1, t1 in REAL_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_${k1}$_${k1}$(x, dim) result(res)
+      module function mean_${rank}$_${k1}$_${k1}$(x, dim, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         integer, intent(in) :: dim
+        logical, intent(in), optional :: mask
         ${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
 
+        if (.not.optval(mask, .true.)) then
+          res = ieee_value(res, ieee_quiet_nan)
+          return
+        end if
+
         if (dim >= 1 .and. dim <= ${rank}$) then
           res = sum(x, dim) / real(size(x, dim), ${k1}$)
         else
@@ -56,13 +76,19 @@ contains
 
   #:for k1, t1 in INT_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_${k1}$_dp(x, dim) result(res)
+      module function mean_${rank}$_${k1}$_dp(x, dim, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         integer, intent(in) :: dim
+        logical, intent(in), optional :: mask
         real(dp) :: res${reduced_shape('x', rank, 'dim')}$
 
+        if (.not.optval(mask, .true.)) then
+          res = ieee_value(res, ieee_quiet_nan)
+          return
+        end if
+
         if (dim >= 1 .and. dim <= ${rank}$) then
-          res = sum(x, dim) / real(size(x, dim), dp)
+          res = sum(real(x, dp), dim) / real(size(x, dim), dp)
         else
           call error_stop("ERROR (mean): wrong dimension")
         end if
@@ -71,4 +97,70 @@ contains
     #:endfor
   #:endfor
 
+
+  #:for k1, t1 in REAL_KINDS_TYPES
+    #:for rank in RANKS
+      module function mean_${rank}$_mask_all_${k1}$_${k1}$(x, mask) result(res)
+        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        logical, intent(in) :: mask${ranksuffix(rank)}$
+        ${t1}$ :: res
+
+        res = sum(x, mask) / real(count(mask, kind = int64), ${k1}$)
+
+      end function mean_${rank}$_mask_all_${k1}$_${k1}$
+    #:endfor
+  #:endfor
+
+
+  #:for k1, t1 in INT_KINDS_TYPES
+    #:for rank in RANKS
+      module function mean_${rank}$_mask_all_${k1}$_dp(x, mask) result(res)
+        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        logical, intent(in) :: mask${ranksuffix(rank)}$
+        real(dp) :: res
+
+        res = sum(real(x, dp), mask) / real(count(mask, kind = int64), dp)
+
+      end function mean_${rank}$_mask_all_${k1}$_dp
+    #:endfor
+  #:endfor
+
+
+  #:for k1, t1 in REAL_KINDS_TYPES
+    #:for rank in RANKS
+      module function mean_${rank}$_mask_${k1}$_${k1}$(x, dim, mask) result(res)
+        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        integer, intent(in) :: dim
+        logical, intent(in) :: mask${ranksuffix(rank)}$
+        ${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
+
+        if (dim >= 1 .and. dim <= ${rank}$) then
+          res = sum(x, dim, mask) / real(count(mask, dim), ${k1}$)
+        else
+          call error_stop("ERROR (mean): wrong dimension")
+        end if
+
+      end function mean_${rank}$_mask_${k1}$_${k1}$
+    #:endfor
+  #:endfor
+
+
+  #:for k1, t1 in INT_KINDS_TYPES
+    #:for rank in RANKS
+      module function mean_${rank}$_mask_${k1}$_dp(x, dim, mask) result(res)
+        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        integer, intent(in) :: dim
+        logical, intent(in) :: mask${ranksuffix(rank)}$
+        real(dp) :: res${reduced_shape('x', rank, 'dim')}$
+
+        if (dim >= 1 .and. dim <= ${rank}$) then
+          res = sum(real(x, dp), dim, mask) / real(count(mask, dim), dp)
+        else
+          call error_stop("ERROR (mean): wrong dimension")
+        end if
+
+      end function mean_${rank}$_mask_${k1}$_dp
+    #:endfor
+  #:endfor
+
 end submodule

src/tests/stats/test_mean.f90

@@ -37,6 +37,17 @@ program test_mean
 call assert( sum( abs( mean(d,2) - sum(d,2)/real(size(d,2), dp) )) < dptol)
 
 
+! check mask = .false.
+
+call assert( isnan(mean(d, .false.)))
+call assert( any(isnan(mean(d, 1, .false.))))
+call assert( any(isnan(mean(d, 2, .false.))))
+
+! check mask of the same shape as input
+call assert( abs(mean(d, d > 0) - sum(d, d > 0)/real(count(d > 0), dp)) < dptol)
+call assert( sum(abs(mean(d, 1, d > 0) - sum(d, 1, d > 0)/real(count(d > 0, 1), dp))) < dptol)
+call assert( sum(abs(mean(d, 2, d > 0) - sum(d, 2, d > 0)/real(count(d > 0, 2), dp))) < dptol)
+
 !int32
 call loadtxt("array3.dat", d)
 
@@ -56,8 +67,8 @@ program test_mean
 !dp rank 3
 allocate(d3(size(d,1),size(d,2),3))
 d3(:,:,1)=d;
-d3(:,:,2)=d*1.5_dp;
-d3(:,:,3)=d*4._dp;
+d3(:,:,2)=d*1.5;
+d3(:,:,3)=d*4;
 
 call assert( abs(mean(d3) - sum(d3)/real(size(d3), dp)) < dptol)
 call assert( sum( abs( mean(d3,1) - sum(d3,1)/real(size(d3,1), dp) )) < dptol)
@@ -67,16 +78,32 @@ program test_mean
 
 !dp rank 4
 allocate(d4(size(d,1),size(d,2),3,9))
-d4 = 1.
+d4 = -1
 d4(:,:,1,1)=d;
-d4(:,:,2,1)=d*1.5_dp;
-d4(:,:,3,1)=d*4._dp;
-d4(:,:,3,9)=d*4._dp;
+d4(:,:,2,1)=d*1.5;
+d4(:,:,3,1)=d*4;
+d4(:,:,3,9)=d*4;
 
 call assert( abs(mean(d4) - sum(d4)/real(size(d4), dp)) < dptol)
 call assert( sum( abs( mean(d4,1) - sum(d4,1)/real(size(d4,1), dp) )) < dptol)
 call assert( sum( abs( mean(d4,2) - sum(d4,2)/real(size(d4,2), dp) )) < dptol)
 call assert( sum( abs( mean(d4,3) - sum(d4,3)/real(size(d4,3), dp) )) < dptol)
 call assert( sum( abs( mean(d4,4) - sum(d4,4)/real(size(d4,4), dp) )) < dptol)
 
+! check mask = .false.
+
+call assert( isnan(mean(d4, .false.)))
+call assert( any(isnan(mean(d4, 1, .false.))))
+call assert( any(isnan(mean(d4, 2, .false.))))
+call assert( any(isnan(mean(d4, 3, .false.))))
+call assert( any(isnan(mean(d4, 4, .false.))))
+
+
+! check mask of the same shape as input
+call assert( abs(mean(d4, d4 > 0) - sum(d4, d4 > 0)/real(count(d4 > 0), dp)) < dptol)
+call assert( any(isnan(mean(d4, 1, d4 > 0))) )
+call assert( any(isnan(mean(d4, 2, d4 > 0))) )
+call assert( any(isnan(mean(d4, 3, d4 > 0))) )
+call assert( sum(abs(mean(d4, 4, d4 > 0) - sum(d4, 4, d4 > 0)/real(count(d4 > 0, 4), dp))) < dptol)
+
 end program