Added support and test for mean of complex arrays

src/common.fypp

@@ -97,5 +97,20 @@ ${prefix + joinstr.join([line.strip() for line in txt.split("\n")]) + suffix}$
   #:endif
 #:enddef
 
+#! Generates a routine name from a generic name, rank, type and kind
+#!
+#! Args:
+#!   gname (str): Generic name
+#!   rank (integer): Rank if exist
+#!   type (str): Type of the input
+#!   kind (str): kind of inputs variable
+#!   suffix (str): other identifier (could be used for output type/kind)
+#!
+#! Returns:
+#!   A string with a new name
+#!
+#:def rname(gname, rank, type, kind, suffix='')
+  $:"{0}_{1}_{2}{3}_{2}{3}".format(gname, rank, type[0], kind) if suffix == '' else "{0}_{1}_{2}{3}_{4}".format(gname, rank, type[0], kind, suffix)
+#:enddef
 
 #:endmute

src/stdlib_experimental_stats.fypp

@@ -1,8 +1,6 @@
 #:include "common.fypp"
-
 #:set RANKS = range(1, MAXRANK + 1)
-
-
+#:set RC_KINDS_TYPES = REAL_KINDS_TYPES + CMPLX_KINDS_TYPES
 module stdlib_experimental_stats
   use stdlib_experimental_kinds, only: sp, dp, qp, &
       int8, int16, int32, int64
@@ -12,92 +10,95 @@ module stdlib_experimental_stats
   public :: mean
 
   interface mean
-
-    #:for k1, t1 in REAL_KINDS_TYPES
+    #:for k1, t1 in RC_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_all_${k1}$_${k1}$(x, mask) result(res)
+        #:set RName = rname("mean_all",rank, t1, k1)
+        module function ${RName}$ (x, mask) result(res)
           ${t1}$, intent(in) :: x${ranksuffix(rank)}$
           logical, intent(in), optional :: mask
           ${t1}$ :: res
-        end function mean_${rank}$_all_${k1}$_${k1}$
+        end function ${RName}$
       #:endfor
     #:endfor
 
     #:for k1, t1 in INT_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_all_${k1}$_dp(x, mask) result(res)
+      #:set RName = rname('mean_all', rank, t1, k1,'dp')
+        module function ${RName}$(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
+        end function ${RName}$
       #:endfor
     #:endfor
 
-    #:for k1, t1 in REAL_KINDS_TYPES
+    #:for k1, t1 in RC_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_${k1}$_${k1}$(x, dim, mask) result(res)
+      #:set RName = rname("mean",rank, t1, k1)
+        module function ${RName}$(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}$
+        end function ${RName}$
       #:endfor
     #:endfor
 
     #:for k1, t1 in INT_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_${k1}$_dp(x, dim, mask) result(res)
+        #:set RName = rname("mean",rank, t1, k1,'dp')
+        module function ${RName}$(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
+        end function ${RName}$
       #:endfor
     #:endfor
 
-
-    #:for k1, t1 in REAL_KINDS_TYPES
+    #:for k1, t1 in RC_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_mask_all_${k1}$_${k1}$(x, mask) result(res)
+        #:set RName = rname('mean_mask_all',rank, t1, k1)
+        module function ${RName}$(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}$
+          end function ${RName}$
       #: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)
+      #:set RName = rname('mean_mask_all',rank, t1, k1, 'dp')
+        module function ${RName}$(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
+        end function ${RName}$
       #:endfor
     #:endfor
 
-
-    #:for k1, t1 in REAL_KINDS_TYPES
+    #:for k1, t1 in RC_KINDS_TYPES
       #:for rank in RANKS
-        module function mean_${rank}$_mask_${k1}$_${k1}$(x, dim, mask) result(res)
+      #:set RName = rname('mean_mask',rank, t1, k1)
+        module function ${RName}$(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}$
+        end function ${RName}$
       #: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)
+      #:set RName = rname('mean_mask',rank, t1, k1, 'dp')
+        module function ${RName}$(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
+        end function ${RName}$
       #:endfor
     #:endfor
 

src/stdlib_experimental_stats.md

@@ -18,15 +18,15 @@ Returns the mean of all the elements of `array`, or of the elements of `array` a
 
 ### Arguments
 
-`array`: Shall be an array of type `integer`, or `real`.
+`array`: Shall be an array of type `integer`, `real`, or `complex`.
 
 `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 `real` or `complex`, 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.

src/stdlib_experimental_stats_mean.fypp

@@ -1,8 +1,6 @@
 #:include "common.fypp"
-
 #:set RANKS = range(1, MAXRANK + 1)
-
-
+#:set RC_KINDS_TYPES = REAL_KINDS_TYPES + CMPLX_KINDS_TYPES
 submodule (stdlib_experimental_stats) stdlib_experimental_stats_mean
 
   use, intrinsic:: ieee_arithmetic, only: ieee_value, ieee_quiet_nan
@@ -12,28 +10,29 @@ submodule (stdlib_experimental_stats) stdlib_experimental_stats_mean
 
 contains
 
-  #:for k1, t1 in REAL_KINDS_TYPES
+  #:for k1, t1 in RC_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_all_${k1}$_${k1}$(x, mask) result(res)
+      #:set RName = rname("mean_all",rank, t1, k1)
+      module function ${RName}$ (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)
+          res = ieee_value(real(res, kind=${k1}$), ieee_quiet_nan)
           return
         end if
 
         res = sum(x) / real(size(x, kind = int64), ${k1}$)
 
-      end function mean_${rank}$_all_${k1}$_${k1}$
+      end function ${RName}$
     #:endfor
   #:endfor
 
-
   #:for k1, t1 in INT_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_all_${k1}$_dp(x, mask) result(res)
+      #:set RName = rname('mean_all', rank, t1, k1,'dp')
+      module function ${RName}$(x, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         logical, intent(in), optional :: mask
         real(dp) :: res
@@ -45,21 +44,22 @@ contains
 
         res = sum(real(x, dp)) / real(size(x, kind = int64), dp)
 
-      end function mean_${rank}$_all_${k1}$_dp
+      end function ${RName}$
     #:endfor
   #:endfor
 
 
-  #:for k1, t1 in REAL_KINDS_TYPES
+  #:for k1, t1 in RC_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_${k1}$_${k1}$(x, dim, mask) result(res)
+      #:set RName = rname("mean",rank, t1, k1)
+      module function ${RName}$(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)
+          res = ieee_value(real(res, kind=${k1}$), ieee_quiet_nan)
           return
         end if
 
@@ -69,14 +69,15 @@ contains
           call error_stop("ERROR (mean): wrong dimension")
         end if
 
-      end function mean_${rank}$_${k1}$_${k1}$
+      end function ${RName}$
     #:endfor
   #:endfor
 
 
   #:for k1, t1 in INT_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_${k1}$_dp(x, dim, mask) result(res)
+      #:set RName = rname("mean",rank, t1, k1,'dp')
+      module function ${RName}$(x, dim, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         integer, intent(in) :: dim
         logical, intent(in), optional :: mask
@@ -93,42 +94,43 @@ contains
           call error_stop("ERROR (mean): wrong dimension")
         end if
 
-      end function mean_${rank}$_${k1}$_dp
+      end function ${RName}$
     #:endfor
   #:endfor
 
-
-  #:for k1, t1 in REAL_KINDS_TYPES
+  #:for k1, t1 in RC_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_mask_all_${k1}$_${k1}$(x, mask) result(res)
+      #:set RName = rname('mean_mask_all',rank, t1, k1)
+      module function ${RName}$(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}$
+      end function ${RName}$
     #: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)
+      #:set RName = rname('mean_mask_all',rank, t1, k1, 'dp')
+      module function ${RName}$(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
+      end function ${RName}$
     #:endfor
   #:endfor
 
-
-  #:for k1, t1 in REAL_KINDS_TYPES
+  #:for k1, t1 in RC_KINDS_TYPES
     #:for rank in RANKS
-      module function mean_${rank}$_mask_${k1}$_${k1}$(x, dim, mask) result(res)
+      #:set RName = rname('mean_mask',rank, t1, k1)
+      module function  ${RName}$(x, dim, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         integer, intent(in) :: dim
         logical, intent(in) :: mask${ranksuffix(rank)}$
@@ -140,14 +142,15 @@ contains
           call error_stop("ERROR (mean): wrong dimension")
         end if
 
-      end function mean_${rank}$_mask_${k1}$_${k1}$
+      end function ${RName}$
     #: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)
+      #:set RName = rname('mean_mask',rank, t1, k1, 'dp')
+      module function ${RName}$(x, dim, mask) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         integer, intent(in) :: dim
         logical, intent(in) :: mask${ranksuffix(rank)}$
@@ -159,7 +162,7 @@ contains
           call error_stop("ERROR (mean): wrong dimension")
         end if
 
-      end function mean_${rank}$_mask_${k1}$_dp
+      end function ${RName}$
     #:endfor
   #:endfor
 

src/tests/stats/test_mean.f90

@@ -13,9 +13,13 @@ program test_mean
 real(sp), allocatable :: s(:, :)
 real(dp), allocatable :: d(:, :)
 
+complex(dp), allocatable :: cs(:, :)
+complex(dp), allocatable :: cd(:, :)
+
 real(dp), allocatable :: d3(:, :, :)
 real(dp), allocatable :: d4(:, :, :, :)
 
+complex(dp), allocatable :: cd3(:, :, :)
 
 !sp
 call loadtxt("array3.dat", s)
@@ -36,6 +40,23 @@ program test_mean
 call assert( sum( abs( mean(d,1) - sum(d,1)/real(size(d,1), dp) )) < dptol)
 call assert( sum( abs( mean(d,2) - sum(d,2)/real(size(d,2), dp) )) < dptol)
 
+!csp
+
+call loadtxt("array3.dat", d)
+cs = cmplx(1._sp, 1._sp)*d
+
+call assert( abs(mean(cs) - sum(cs)/real(size(cs), sp)) < sptol)
+call assert( sum( abs( mean(cs,1) - sum(cs,1)/real(size(cs,1), sp) )) < sptol)
+call assert( sum( abs( mean(cs,2) - sum(cs,2)/real(size(cs,2), sp) )) < sptol)
+
+!cdp
+
+call loadtxt("array3.dat", d)
+cd = cmplx(1._dp, 1._dp)*d
+
+call assert( abs(mean(cd) - sum(cd)/real(size(cd), dp)) < dptol)
+call assert( sum( abs( mean(cd,1) - sum(cd,1)/real(size(cd,1), dp) )) < dptol)
+call assert( sum( abs( mean(cd,2) - sum(cd,2)/real(size(cd,2), dp) )) < dptol)
 
 ! check mask = .false.
 
@@ -75,6 +96,18 @@ program test_mean
 call assert( sum( abs( mean(d3,2) - sum(d3,2)/real(size(d3,2), dp) )) < dptol)
 call assert( sum( abs( mean(d3,3) - sum(d3,3)/real(size(d3,3), dp) )) < dptol)
 
+!cdp rank 3
+allocate(cd3(size(d,1),size(d,2),3))
+cd3(:,:,1)=d;
+cd3(:,:,2)=d*1.5;
+cd3(:,:,3)=d*4;
+cd3 = cmplx(1._sp, 1._sp)*cd3
+
+call assert( abs(mean(cd3) - sum(cd3)/real(size(cd3), dp)) < dptol)
+call assert( sum( abs( mean(cd3,1) - sum(cd3,1)/real(size(cd3,1), dp) )) < dptol)
+call assert( sum( abs( mean(cd3,2) - sum(cd3,2)/real(size(cd3,2), dp) )) < dptol)
+call assert( sum( abs( mean(cd3,3) - sum(cd3,3)/real(size(cd3,3), dp) )) < dptol)
+
 
 !dp rank 4
 allocate(d4(size(d,1),size(d,2),3,9))