fortran-lang · certik · Jan 29, 2020 · Jan 29, 2020 · Jan 29, 2020
diff --git a/src/common.fypp b/src/common.fypp
@@ -37,7 +37,7 @@
 #! Generates an array rank suffix.
 #!
 #! Args:
-#!     rank [in]: Integer with rank.
+#!     rank (int): Rank of the variable
 #!
 #! Returns:
 #!     Array rank suffix string (e.g. (:,:) if rank = 2)
@@ -46,4 +46,48 @@
 #{if rank > 0}#(${":" + ",:" * (rank - 1)}$)#{endif}#
 #:enddef
 
+
+#! Joins stripped lines with given character string
+#!
+#! Args:
+#!   txt (str): Text to process
+#!   joinstr (str): String to use as connector
+#!   prefix (str): String to add as prefix before the joined text
+#!   suffix (str): String to add as suffix after the joined text
+#!
+#! Returns:
+#!   Lines stripped and joined with the given string.
+#!
+#:def join_lines(txt, joinstr, prefix="", suffix="")
+${prefix + joinstr.join([line.strip() for line in txt.split("\n")]) + suffix}$
+#:enddef
+
+
+#! Brace enclosed, comma separated Fortran expressions for a reduced shape.
+#!
+#! Rank of the original variable will be reduced by one. The routine generates
+#! for each dimension a Fortan expression using merge(), which calculates the
+#! size of the array for that dimension.
+#!
+#! Args:
+#!   varname (str): Name of the variable to be used as origin
+#!   origrank (int): Rank of the original variable
+#!   idim (int): Index of the reduced dimension
+#!
+#! Returns:
+#!   Shape expression enclosed in braces, so that it can be used as suffix to
+#!   define array shapes in declarations.
+#!
+#:def reduced_shape(varname, origrank, idim)
+  #:assert origrank > 0
+  #:if origrank > 1
+    #:call join_lines(joinstr=", ", prefix="(", suffix=")")
+      #:for i in range(1, origrank)
+        merge(size(${varname}$, ${i}$), size(${varname}$, ${i + 1}$), mask=${i}$<${idim}$)
+      #:endfor
+    #:endcall
+  #:endif
+#:enddef
+
+
 #:endmute
diff --git a/src/stdlib_experimental_stats.fypp b/src/stdlib_experimental_stats.fypp
@@ -1,6 +1,6 @@
 #:include "common.fypp"
 
-#:set RANKS = range(3, MAXRANK + 1)
+#:set RANKS = range(1, MAXRANK + 1)
 
 
 module stdlib_experimental_stats
@@ -12,49 +12,6 @@ module stdlib_experimental_stats
   public :: mean
 
   interface mean
-    #:for k1, t1 in REAL_KINDS_TYPES
-      module function mean_1_${k1}$_${k1}$(x) result(res)
-        ${t1}$, intent(in) :: x(:)
-        ${t1}$ :: res
-      end function mean_1_${k1}$_${k1}$
-    #:endfor
-
-    #:for k1, t1 in INT_KINDS_TYPES
-      module function mean_1_${k1}$_dp(x) result(res)
-        ${t1}$, intent(in) :: x(:)
-        real(dp) :: res
-      end function mean_1_${k1}$_dp
-    #:endfor
-
-    #:for k1, t1 in REAL_KINDS_TYPES
-      module function mean_2_all_${k1}$_${k1}$(x) result(res)
-          ${t1}$, intent(in) :: x(:,:)
-          ${t1}$ :: res
-        end function mean_2_all_${k1}$_${k1}$
-    #:endfor
-
-    #:for k1, t1 in INT_KINDS_TYPES
-      module function mean_2_all_${k1}$_dp(x) result(res)
-        ${t1}$, intent(in) :: x(:,:)
-        real(dp) :: res
-      end function mean_2_all_${k1}$_dp
-    #:endfor
-
-    #:for k1, t1 in REAL_KINDS_TYPES
-      module function mean_2_${k1}$_${k1}$(x, dim) result(res)
-        ${t1}$, intent(in) :: x(:,:)
-        integer, intent(in) :: dim
-        ${t1}$ :: res(merge(size(x, 1), size(x, 2), mask = 1 < dim ))
-      end function mean_2_${k1}$_${k1}$
-    #:endfor
-
-    #:for k1, t1 in INT_KINDS_TYPES
-      module function mean_2_${k1}$_dp(x, dim) result(res)
-        ${t1}$, intent(in) :: x(:,:)
-        integer, intent(in) :: dim
-        real(dp) :: res(merge(size(x, 1), size(x, 2), mask = 1 < dim ))
-      end function mean_2_${k1}$_dp
-    #:endfor
 
     #:for k1, t1 in REAL_KINDS_TYPES
       #:for rank in RANKS
@@ -76,34 +33,23 @@ module stdlib_experimental_stats
 
     #:for k1, t1 in REAL_KINDS_TYPES
       #:for rank in RANKS
-      module function mean_${rank}$_${k1}$_${k1}$(x, dim) result(res)
-        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
-        integer, intent(in) :: dim
-        ${t1}$ :: res( &
-        #:for imerge in range(1,rank-1)
-            & merge(size(x, ${imerge}$),size(x, ${imerge + 1}$),&
-            & mask = ${imerge}$ < dim), &
-        #:endfor
-          & merge(size(x, ${rank-1}$), size(x, ${rank}$),&
-          & mask = ${rank-1}$ < dim))
-      end function mean_${rank}$_${k1}$_${k1}$
+        module function mean_${rank}$_${k1}$_${k1}$(x, dim) result(res)
+          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          integer, intent(in) :: dim
+          ${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
+        end function mean_${rank}$_${k1}$_${k1}$
+      #:endfor
     #:endfor
-  #:endfor
 
-  #:for k1, t1 in INT_KINDS_TYPES
-    #:for rank in RANKS
-      module function mean_${rank}$_${k1}$_dp(x, dim) result(res)
-        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
-        integer, intent(in) :: dim
-        real(dp) :: res( &
-        #:for imerge in range(1,rank-1)
-            & merge(size(x, ${imerge}$), size(x,${imerge + 1}$),&
-            & mask = ${imerge}$ < dim), &
-        #:endfor
-          & merge(size(x,${rank-1}$),size(x,${rank}$),mask = ${rank-1}$ < dim ))
-      end function mean_${rank}$_${k1}$_dp
+    #:for k1, t1 in INT_KINDS_TYPES
+      #:for rank in RANKS
+        module function mean_${rank}$_${k1}$_dp(x, dim) result(res)
+          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+          integer, intent(in) :: dim
+          real(dp) :: res${reduced_shape('x', rank, 'dim')}$
+        end function mean_${rank}$_${k1}$_dp
+      #:endfor
     #:endfor
-  #:endfor
 
   end interface mean
 

diff --git a/src/stdlib_experimental_stats_mean.fypp b/src/stdlib_experimental_stats_mean.fypp
@@ -1,6 +1,6 @@
 #:include "common.fypp"
 
-#:set RANKS = range(3, MAXRANK + 1)
+#:set RANKS = range(1, MAXRANK + 1)
 
 
 submodule (stdlib_experimental_stats) stdlib_experimental_stats_mean
@@ -10,88 +10,6 @@ submodule (stdlib_experimental_stats) stdlib_experimental_stats_mean
 
 contains
 
-  #:for k1, t1 in REAL_KINDS_TYPES
-    module function mean_1_${k1}$_${k1}$(x) result(res)
-      ${t1}$, intent(in) :: x(:)
-      ${t1}$ :: res
-
-      res = sum(x) / real(size(x, kind = int64), ${k1}$)
-
-    end function mean_1_${k1}$_${k1}$
-  #:endfor
-
-
-  #:for k1, t1 in INT_KINDS_TYPES
-    module function mean_1_${k1}$_dp(x) result(res)
-      ${t1}$, intent(in) :: x(:)
-      real(dp) :: res
-
-      res = sum(real(x, dp)) / real(size(x, kind = int64), dp)
-
-    end function mean_1_${k1}$_dp
-  #:endfor
-
-
-  #:for k1, t1 in REAL_KINDS_TYPES
-    module function mean_2_all_${k1}$_${k1}$(x) result(res)
-      ${t1}$, intent(in) :: x(:,:)
-      ${t1}$ :: res
-
-      res = sum(x) / real(size(x, kind = int64), ${k1}$)
-
-    end function mean_2_all_${k1}$_${k1}$
-  #:endfor
-
-
-  #:for k1, t1 in INT_KINDS_TYPES
-    module function mean_2_all_${k1}$_dp(x) result(res)
-      ${t1}$, intent(in) :: x(:,:)
-      real(dp) :: res
-
-      res = sum(real(x, dp)) / real(size(x, kind = int64), dp)
-
-    end function mean_2_all_${k1}$_dp
-  #:endfor
-
-
-  #:for k1, t1 in REAL_KINDS_TYPES
-    module function mean_2_${k1}$_${k1}$(x, dim) result(res)
-      ${t1}$, intent(in) :: x(:,:)
-      integer, intent(in) :: dim
-      ${t1}$ :: res(merge(size(x, 1), size(x, 2), mask = 1 < dim ))
-
-      select case(dim)
-      case(1)
-        res = sum(x, 1) / real(size(x, 1), ${k1}$)
-      case(2)
-        res = sum(x, 2) / real(size(x, 2), ${k1}$)
-      case default
-        call error_stop("ERROR (mean): wrong dimension")
-      end select
-
-    end function mean_2_${k1}$_${k1}$
-  #:endfor
-
-
-  #:for k1, t1 in INT_KINDS_TYPES
-    module function mean_2_${k1}$_dp(x, dim) result(res)
-      ${t1}$, intent(in) :: x(:,:)
-      integer, intent(in) :: dim
-      real(dp) :: res(merge(size(x, 1), size(x, 2), mask = 1 < dim ))
-
-      select case(dim)
-      case(1)
-        res = sum(real(x, dp), 1) / real(size(x, 1), dp)
-      case(2)
-        res = sum(real(x, dp), 2) / real(size(x, 2), dp)
-      case default
-        call error_stop("ERROR (mean): wrong dimension")
-      end select
-
-    end function mean_2_${k1}$_dp
-  #:endfor
-
-
   #:for k1, t1 in REAL_KINDS_TYPES
     #:for rank in RANKS
       module function mean_${rank}$_all_${k1}$_${k1}$(x) result(res)
@@ -123,22 +41,13 @@ contains
       module function mean_${rank}$_${k1}$_${k1}$(x, dim) result(res)
         ${t1}$, intent(in) :: x${ranksuffix(rank)}$
         integer, intent(in) :: dim
-        ${t1}$ :: res( &
-          #:for imerge in range(1,rank-1)
-            merge(size(x,${imerge}$),size(x,${imerge + 1}$),&
-              & mask = ${imerge}$ < dim ), &
-          #:endfor
-            & merge(size(x,${rank-1}$),size(x,${rank}$),&
-            & mask = ${rank-1}$ < dim ))
-
-        select case(dim)
-        #:for fi in range(1,rank+1)
-          case(${fi}$)
-            res=sum(x, ${fi}$) / real(size(x, ${fi}$), ${k1}$)
-        #:endfor
-        case default
+        ${t1}$ :: res${reduced_shape('x', rank, 'dim')}$
+
+        if (dim >= 1 .and. dim <= ${rank}$) then
+          res = sum(x, dim) / real(size(x, dim), ${k1}$)
+        else
           call error_stop("ERROR (mean): wrong dimension")
-        end select
+        end if
 
       end function mean_${rank}$_${k1}$_${k1}$
     #:endfor
@@ -148,24 +57,15 @@ contains
   #:for k1, t1 in INT_KINDS_TYPES
     #:for rank in RANKS
       module function mean_${rank}$_${k1}$_dp(x, dim) result(res)
-          ${t1}$, intent(in) :: x${ranksuffix(rank)}$
-          integer, intent(in) :: dim
-          real(dp) :: res( &
-          #:for imerge in range(1,rank-1)
-              & merge(size(x, ${imerge}$), size(x, ${imerge + 1}$),&
-              & mask = ${imerge}$ < dim ), &
-          #:endfor
-              & merge(size(x,${rank-1}$),size(x,${rank}$),&
-              & mask = ${rank-1}$ < dim ))
-
-          select case(dim)
-           #:for fi in range(1,rank+1)
-             case(${fi}$)
-               res = sum(real(x, dp), ${fi}$) / real(size(x, ${fi}$), dp)
-           #:endfor
-           case default
-             call error_stop("ERROR (mean): wrong dimension")
-          end select
+        ${t1}$, intent(in) :: x${ranksuffix(rank)}$
+        integer, intent(in) :: dim
+        real(dp) :: res${reduced_shape('x', rank, 'dim')}$
+
+        if (dim >= 1 .and. dim <= ${rank}$) then
+          res = sum(x, dim) / real(size(x, dim), dp)
+        else
+          call error_stop("ERROR (mean): wrong dimension")
+        end if
 
       end function mean_${rank}$_${k1}$_dp
     #:endfor