median: return NaN when real array contain NaN

Author: jvdp1

doc/specs/stdlib_stats.md

@@ -229,12 +229,14 @@ Generic subroutine
 ### Return value
 
 If `array` is of type `real`, the result is of type `real` with the same kind as `array`.
+If `array` is of type `real` and contains IEEE `NaN`, the result is IEEE `NaN`.
 If `array` is of type `integer`, the result is of type `real(dp)`.
 
 If `dim` is absent, a scalar with the median 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 median 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

src/stdlib_stats_median.fypp

@@ -4,7 +4,7 @@
 
 submodule (stdlib_stats) stdlib_stats_median
 
-  use, intrinsic:: ieee_arithmetic, only: ieee_value, ieee_quiet_nan
+  use, intrinsic:: ieee_arithmetic, only: ieee_value, ieee_quiet_nan, ieee_is_nan
   use stdlib_error, only: error_stop
   use stdlib_optval, only: optval
   ! Use "ord_sort" rather than "sort" because the former can be much faster for arrays
@@ -31,6 +31,13 @@ contains
           return
         end if
 
+        #:if t1[0] == 'r'
+          if (any(ieee_is_nan(x))) then
+            res = ieee_value(1._${o1}$, ieee_quiet_nan)
+            return
+          end if
+        #:endif
+
         n = size(x, kind=int64)
         c = floor( (n + 1) / 2._${o1}$, kind=int64 )
 
@@ -87,6 +94,19 @@ contains
               do j${fj}$ = 1, size(x, ${fj}$)
             #:endfor
                 x_tmp(:) = x${select_subvector('j', rank, fi)}$
+
+                #:if t1[0] == 'r'
+                  if (any(ieee_is_nan(x_tmp))) then
+                    res${reduce_subvector('j', rank, fi)}$ = &
+                      ieee_value(1._${o1}$, ieee_quiet_nan)
+                    #:if fi == 1
+                      return
+                    #:else
+                      cycle
+                    #:endif
+                  end if
+                #:endif
+
                 call sort(x_tmp)
 
                 if (mod(n, 2) == 0) then
@@ -127,12 +147,19 @@ contains
             call error_stop("ERROR (median): shapes of x and mask are different")
         end if
 
+        #:if t1[0] == 'r'
+          if (any(ieee_is_nan(x))) then
+            res = ieee_value(1._${o1}$, ieee_quiet_nan)
+            return
+          end if
+        #:endif
+
         x_tmp = pack(x, mask)
 
         call sort(x_tmp)
 
-        n = size(x_tmp, kind=int64 )
-        c = floor( (n + 1) / 2._${o1}$, kind=int64 )
+        n = size(x_tmp, kind=int64)
+        c = floor( (n + 1) / 2._${o1}$, kind=int64)
 
         if (n == 0) then
             res = ieee_value(1._${o1}$, ieee_quiet_nan)
@@ -180,6 +207,19 @@ contains
             #:endfor
                 x_tmp = pack(x${select_subvector('j', rank, fi)}$, &
                               mask${select_subvector('j', rank, fi)}$)
+
+                #:if t1[0] == 'r'
+                  if (any(ieee_is_nan(x_tmp))) then
+                    res${reduce_subvector('j', rank, fi)}$ = &
+                      ieee_value(1._${o1}$, ieee_quiet_nan)
+                    #:if rank == 1
+                      return
+                    #:else
+                      cycle
+                    #:endif
+                  end if
+                #:endif
+
                 call sort(x_tmp)
 
                 n = size(x_tmp, kind=int64)

src/tests/stats/test_median.fypp

@@ -5,7 +5,7 @@ program test_median
     use stdlib_kinds, only: sp, dp, qp, &
         int8, int16, int32, int64
     use stdlib_stats, only: median, mean
-    use, intrinsic :: ieee_arithmetic, only : ieee_is_nan
+    use, intrinsic :: ieee_arithmetic, only : ieee_is_nan, ieee_value, ieee_quiet_nan
     implicit none
     real(sp), parameter :: sptol = 1000 * epsilon(1._sp)
     real(dp), parameter :: dptol = 2000 * epsilon(1._dp)
@@ -108,6 +108,7 @@ end subroutine
 #:for k1,t1 in REAL_KINDS_TYPES
 subroutine test_median_${k1}$()
     ${t1}$, allocatable :: d0(:), d1(:), d2(:,:), d3(:,:,:)
+    ${t1}$, allocatable :: tmp1(:)
 
     allocate(d0(0))
     !check just to be sure that the setup of d0 is correct
@@ -190,6 +191,41 @@ subroutine test_median_${k1}$()
     call check( sum(abs(median(d2, 2) - [7._${k1}$, 1._${k1}$, 0._${k1}$])) < ${k1}$tol,                &
                  '${k1}$ median(d2, 2), odd')
 
+    !check IEEE NaN values in array
+    d1(1) = ieee_value(1._${k1}$, ieee_quiet_nan)
+    d2(1, 1) =  ieee_value(1._${k1}$, ieee_quiet_nan)
+    d3(1, 1, 1) =  ieee_value(1._${k1}$, ieee_quiet_nan)
+
+    !median_all
+    call check( ieee_is_nan(median(d1)), '${k1}$ median(d1), should be NaN')
+    call check( ieee_is_nan(median(d2)), '${k1}$ median(d2), should be NaN')
+    call check( ieee_is_nan(median(d3)), '${k1}$ median(d3), should be NaN')
+
+    !median
+    call check( any(ieee_is_nan(median(d2, 1))), '${k1}$ median(d2, 1) should contain at least 1 NaN' )
+    call check( any(ieee_is_nan(median(d2, 2))), '${k1}$ median(d2, 2) should contain at least 1 NaN' )
+
+    call check( any(ieee_is_nan(median(d3, 1))), '${k1}$ median(d3, 1) should contain at least 1 NaN' )
+    call check( any(ieee_is_nan(median(d3, 2))), '${k1}$ median(d3, 2) should contain at least 1 NaN' )
+    call check( any(ieee_is_nan(median(d3, 3))), '${k1}$ median(d3, 3) should contain at least 1 NaN' )
+
+    !median_mask_all
+    call check( ieee_is_nan(median(d1, d1 > 0)), '${k1}$ median(d1, d1 > 0) should be NaN' )
+    call check( ieee_is_nan(median(d2, d2 > 0)), '${k1}$ median(d2, d2 > 0) should be NaN' )
+    call check( ieee_is_nan(median(d3, d3 > 0)), '${k1}$ median(d3, d3 > 0) should be NaN' )
+
+    !median mask
+    call check( ieee_is_nan(median(d1, 1, d1 .ne. 0)), '${k1}$ median(d1, 1, d1.ne.0) should return NaN')
+
+    tmp1 = median(d2, 1, d2 .ne. 0)
+    call check( any(ieee_is_nan(tmp1)), '${k1}$ median(d2, 1, d2 .ne. 0 ) should contain at least 1 NaN')
+    call check( tmp1(2) == -4._${k1}$, '${k1}$ tmp1(2) == -4')
+    call check( any(ieee_is_nan(median(d2, 2, d2 .ne. 0))),              &
+       '${k1}$ median(d2, 2, d2 .ne. 0 ) should contain at least 1 NaN')
+
+    call check( any(ieee_is_nan(median(d3, 1, d3 .ne. 0))),              &
+       '${k1}$ median(d3, 1, d3 .ne. 0 ) should contain at least 1 NaN')
+
 end subroutine
 #:endfor