Add equal_nan argument,

zoziha · zoziha · commit adfb820a56cd · 2021-11-10T11:51:21.000+08:00
for `**_close` in stdlib_math.
diff --git a/doc/specs/stdlib_math.md b/doc/specs/stdlib_math.md
@@ -406,7 +406,7 @@ is_close(a, b, rel_tol, abs_tol) = is_close(a%re, b%re, rel_tol, abs_tol) .and.
 
 #### Syntax
 
-`bool = [[stdlib_math(module):is_close(interface)]] (a, b [, rel_tol, abs_tol])`
+`bool = [[stdlib_math(module):is_close(interface)]] (a, b [, rel_tol, abs_tol, equal_nan])`
 
 #### Status
 
@@ -424,12 +424,15 @@ This argument is `intent(in)`.
 `b`: Shall be a `real/complex` scalar/array.
 This argument is `intent(in)`.
 
-`rel_tol`: Shall be a `real` scalar.
+`rel_tol`: Shall be a `real` scalar/array.
 This argument is `intent(in)` and `optional`, which is `1.0e-9` by default.
 
-`abs_tol`: Shall be a `real` scalar.
+`abs_tol`: Shall be a `real` scalar/array.
 This argument is `intent(in)` and `optional`, which is `0.0` by default.
 
+`equal_nan`: Shall be a `logical` scalar/array.
+This argument is `intent(in)` and `optional`, which is `.false.` by default.
+
 Note: All `real/complex` arguments must have same `kind`.  
 If the value of `rel_tol/abs_tol` is negative (not recommended), 
 it will be corrected to `abs(rel_tol/abs_tol)` by the internal process of `is_close`.
@@ -442,25 +445,34 @@ Returns a `logical` scalar/array.
 
 ```fortran
 program demo_math_is_close
+
     use stdlib_math,  only: is_close
     use stdlib_error, only: check
-    real :: x(2) = [1, 2]
-    print *, is_close(x,[real :: 1, 2.1])        !! [T, F]
-    print *, is_close(2.0, 2.1, abs_tol=0.1)     !! T
+    real :: x(2) = [1, 2], y, NAN
+    
+    y   = -3
+    NAN = sqrt(y)
+    
+    print *, is_close(x,[real :: 1, 2.1])       !! [T, F]
+    print *, is_close(2.0, 2.1, abs_tol=0.1)    !! T
+    print *, NAN, is_close(2.0, NAN), is_close(2.0, NAN, equal_nan=.true.)   !! NAN, F, F
+    print *, is_close(NAN, NAN), is_close(NAN, NAN, equal_nan=.true.)        !! F, T
+    
     call check(all(is_close(x, [2.0, 2.0])), msg="all(is_close(x, [2.0, 2.0])) failed.", warn=.true.)
-        !! all(is_close(x, [2.0, 2.0])) failed.
+            !! all(is_close(x, [2.0, 2.0])) failed.
+        
 end program demo_math_is_close
 ```
 
 ### `all_close`
 
 #### Description
 
-Returns a boolean scalar where two arrays are element-wise equal within a tolerance, behaves like `all(is_close(a, b [, rel_tol, abs_tol]))`.
+Returns a boolean scalar where two arrays are element-wise equal within a tolerance, behaves like `all(is_close(a, b [, rel_tol, abs_tol, equal_nan]))`.
 
 #### Syntax
 
-`bool = [[stdlib_math(module):all_close(interface)]] (a, b [, rel_tol, abs_tol])`
+`bool = [[stdlib_math(module):all_close(interface)]] (a, b [, rel_tol, abs_tol, equal_nan])`
 
 #### Status
 
@@ -484,6 +496,9 @@ This argument is `intent(in)` and `optional`, which is `1.0e-9` by default.
 `abs_tol`: Shall be a `real` scalar.
 This argument is `intent(in)` and `optional`, which is `0.0` by default.
 
+`equal_nan`: Shall be a `logical` scalar.
+This argument is `intent(in)` and `optional`, which is `.false.` by default.
+
 Note: All `real/complex` arguments must have same `kind`.  
 If the value of `rel_tol/abs_tol` is negative (not recommended), 
 it will be corrected to `abs(rel_tol/abs_tol)` by the internal process of `all_close`.
@@ -496,18 +511,23 @@ Returns a `logical` scalar.
 
 ```fortran
 program demo_math_all_close
+
     use stdlib_math,  only: all_close
     use stdlib_error, only: check
-    real    :: x(2) = [1, 2], random(4, 4)
+    real    :: x(2) = [1, 2], y, NAN
     complex :: z(4, 4)
     
+    y   = -3
+    NAN = sqrt(y)
+    z   = (1.0, 1.0)
+    
+    print *, all_close(z+cmplx(1.0e-11, 1.0e-11), z)     !! T
+    print *, NAN, all_close([NAN], [NAN]), all_close([NAN], [NAN], equal_nan=.true.) 
+                                                         !! NAN, F, T
+                                                
     call check(all_close(x, [2.0, 2.0], rel_tol=1.0e-6, abs_tol=1.0e-3), &
                msg="all_close(x, [2.0, 2.0]) failed.", warn=.true.)
                !! all_close(x, [2.0, 2.0]) failed.
-
-    call random_number(random(4, 4))
-    z = 1.0
-    print *, all_close(z+1.0e-11*random, z)     !! T
     
 end program demo_math_all_close
 ```
diff --git a/src/stdlib_math.fypp b/src/stdlib_math.fypp
@@ -296,10 +296,10 @@ module stdlib_math
     interface is_close
         #:set RC_KINDS_TYPES = REAL_KINDS_TYPES + CMPLX_KINDS_TYPES
         #:for k1, t1 in RC_KINDS_TYPES
-        elemental module function is_close_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol) result(close)
+        elemental module logical function is_close_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol, equal_nan) result(close)
             ${t1}$, intent(in) :: a, b
             real(${k1}$), intent(in), optional :: rel_tol, abs_tol
-            logical :: close
+            logical, intent(in), optional :: equal_nan
         end function is_close_${t1[0]}$${k1}$
         #:endfor
     end interface is_close
@@ -313,9 +313,10 @@ module stdlib_math
         #:set RANKS = range(1, MAXRANK + 1)
         #:for k1, t1 in RC_KINDS_TYPES
         #:for r1 in RANKS
-        logical pure module function all_close_${r1}$_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol) result(close)
+        logical pure module function all_close_${r1}$_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol, equal_nan) result(close)
             ${t1}$, intent(in) :: a${ranksuffix(r1)}$, b${ranksuffix(r1)}$
             real(${k1}$), intent(in), optional :: rel_tol, abs_tol
+            logical, intent(in), optional :: equal_nan
         end function all_close_${r1}$_${t1[0]}$${k1}$
         #:endfor
         #:endfor
diff --git a/src/stdlib_math_all_close.fypp b/src/stdlib_math_all_close.fypp
@@ -10,12 +10,13 @@ contains
 
     #:for k1, t1 in RC_KINDS_TYPES
     #:for r1 in RANKS
-    logical pure module function all_close_${r1}$_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol) result(close)
+    logical pure module function all_close_${r1}$_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol, equal_nan) result(close)
 
         ${t1}$, intent(in) :: a${ranksuffix(r1)}$, b${ranksuffix(r1)}$
         real(${k1}$), intent(in), optional :: rel_tol, abs_tol
+        logical, intent(in), optional :: equal_nan
 
-        close = all(is_close(a, b, rel_tol, abs_tol))
+        close = all(is_close(a, b, rel_tol, abs_tol, equal_nan))
 
     end function all_close_${r1}$_${t1[0]}$${k1}$
     #:endfor
diff --git a/src/stdlib_math_is_close.fypp b/src/stdlib_math_is_close.fypp
@@ -2,30 +2,40 @@
 
 submodule(stdlib_math) stdlib_math_is_close
 
+    use, intrinsic :: ieee_arithmetic, only: ieee_is_nan
+    implicit none
+    
 contains
 
     #! Determines whether the values of `a` and `b` are close.
 
     #:for k1, t1 in REAL_KINDS_TYPES
-    elemental module function is_close_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol) result(close)
+    elemental module logical function is_close_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol, equal_nan) result(close)
         ${t1}$, intent(in) :: a, b
         real(${k1}$), intent(in), optional :: rel_tol, abs_tol
-        logical :: close
-
-        close = abs(a - b) <= max( abs(optval(rel_tol, 1.0e-9_${k1}$)*max(abs(a), abs(b))), &
-                                   abs(optval(abs_tol, 0.0_${k1}$)) )
+        logical, intent(in), optional :: equal_nan
+        logical :: equal_nan_
+
+        equal_nan_ = optval(equal_nan, .false.)
+        
+        if (ieee_is_nan(a) .or. ieee_is_nan(b)) then
+            close = merge(.true., .false., equal_nan_ .and. ieee_is_nan(a) .and. ieee_is_nan(b))
+        else
+            close = abs(a - b) <= max( abs(optval(rel_tol, 1.0e-9_${k1}$)*max(abs(a), abs(b))), &
+                                       abs(optval(abs_tol, 0.0_${k1}$)) )
+        end if     
 
     end function is_close_${t1[0]}$${k1}$
     #:endfor
 
     #:for k1, t1 in CMPLX_KINDS_TYPES
-    elemental module function is_close_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol) result(close)
+    elemental module logical function is_close_${t1[0]}$${k1}$(a, b, rel_tol, abs_tol, equal_nan) result(close)
         ${t1}$, intent(in) :: a, b
         real(${k1}$), intent(in), optional :: rel_tol, abs_tol
-        logical :: close
+        logical, intent(in), optional :: equal_nan
 
-        close = is_close_r${k1}$(a%re, b%re, rel_tol, abs_tol) .and. &
-                is_close_r${k1}$(a%im, b%im, rel_tol, abs_tol)
+        close = is_close_r${k1}$(a%re, b%re, rel_tol, abs_tol, equal_nan) .and. &
+                is_close_r${k1}$(a%im, b%im, rel_tol, abs_tol, equal_nan)
 
     end function is_close_${t1[0]}$${k1}$
     #:endfor
diff --git a/src/tests/math/test_math_all_close.f90 b/src/tests/math/test_math_all_close.f90
@@ -3,7 +3,11 @@ program tester
     use stdlib_math, only: all_close
     use stdlib_error, only: check
     implicit none
-    
+    real :: y, NAN
+
+    y = -3
+    NAN = sqrt(y)
+
     call test_math_all_close_real
     call test_math_all_close_complex
     print *, "All tests in `test_math_all_close` passed."
@@ -12,26 +16,29 @@ program tester
 
     subroutine test_math_all_close_real
 
-        real :: x(4, 4), random(4, 4)
+        real :: x(4, 4) = 1.0
 
-        call random_number(random)
-        x = 1.0
+        call check(all_close(x + 1.0e-11, x), msg="REAL: all_close(x+1.0e-11, x) failed.")
+        call check(all_close(x + 1.0e-5, x), msg="REAL: all_close(x+1.0e-5 , x) failed. (expected)", warn=.true.)
 
-        call check(all_close(x+1.0e-11*random, x), msg="REAL: all_close(x+1.0e-11*random, x) failed.")
-        call check(all_close(x+1.0e-5 *random, x), msg="REAL: all_close(x+1.0e-5 *random, x) failed.", warn=.true.)
+        !> Tests for NAN
+        call check(all_close(x + NAN, x), msg="REAL: all_close(x+NAN, x) failed. (expected)", warn=.true.)
+        call check(all_close(x + NAN, x + NAN, equal_nan=.true.), msg="REAL: all_close(x+NAN, x, equal_nan=.true.) failed.")
 
     end subroutine test_math_all_close_real
 
     subroutine test_math_all_close_complex
 
-        real    :: random(4, 4)
-        complex :: x(4, 4)
+        complex :: x(4, 4) = cmplx(1.0, 1.0)
 
-        call random_number(random)
-        x = 1.0
+        call check(all_close(x + cmplx((1.0e-15, 1.0e-15)), x), msg="CMPLX: all_close(x+cmplx(1.0e-11, 1.0e-11), x)")
+        call check(all_close(x + cmplx(1.0e-5, 1.0e-5), x), &
+                   msg="CMPLX: all_close(x+cmplx(1.0e-5 , 1.0e-5 ), x) failed. (expected)", warn=.true.)
 
-        call check(all_close(x+1.0e-11*random, x), msg="CMPLX: all_close(x+1.0e-11*random, x)")
-        call check(all_close(x+1.0e-5 *random, x), msg="CMPLX: all_close(x+1.0e-5 *random, x) failed.", warn=.true.)
+        !> Tests for NAN
+        call check(all_close(x + cmplx(NAN, NAN), x), msg="REAL: all_close(x+cmplx(NAN, NAN), x) failed. (expected)", warn=.true.)
+        call check(all_close(x + cmplx(NAN, NAN), x + cmplx(NAN, NAN), equal_nan=.true.), &
+                   msg="REAL: all_close(x+cmplx(NAN, NAN), x, equal_nan=.true.) failed.")
 
     end subroutine test_math_all_close_complex
 
diff --git a/src/tests/math/test_math_is_close.f90 b/src/tests/math/test_math_is_close.f90
@@ -1,5 +1,11 @@
 program test_math_is_close
 
+    implicit none
+
+    real :: x, NAN
+    x   = -3
+    NAN = sqrt(x)
+
     call test_math_is_close_real
     call test_math_is_close_complex
     print *, "All tests in `test_math_is_close` passed."
@@ -12,7 +18,7 @@ subroutine test_math_is_close_real
 
         call check(is_close(2.5, 2.5, rel_tol=1.0e-5), msg="is_close(2.5, 2.5, rel_tol=1.0e-5) failed.")
         call check(all(is_close([2.5, 3.2], [2.5, 10.0], rel_tol=1.0e-5)), &
-                   msg="all(is_close([2.5, 3.2], [2.5, 10.0], rel_tol=1.0e-5)) failed (expected).", warn=.true.)
+                   msg="all(is_close([2.5, 3.2], [2.5, 10.0], rel_tol=1.0e-5)) failed. (expected)", warn=.true.)
         call check(all(is_close(reshape([2.5, 3.2, 2.2, 1.0], [2, 2]), reshape([2.5, 3.2001, 2.25, 1.1], [2, 2]), &
                    abs_tol=1.0e-5, rel_tol=0.1)), &
                    msg="all(is_close(reshape([2.5, 3.2, 2.2, 1.0],[2,2]), reshape([2.5, 3.2001, 2.25, 1.1],[2,2]), &
@@ -21,26 +27,41 @@ subroutine test_math_is_close_real
         !> Tests for zeros
         call check(is_close(0.0, -0.0), msg="is_close(0.0, -0.0) failed.")
 
+        !> Tests for NaN
+        call check(is_close(NAN, NAN), msg="is_close(NAN, NAN) failed.", warn=.true.)
+        call check(is_close(NAN, NAN, equal_nan=.true.), msg="is_close(NAN, NAN, equal_nan=.true.) failed.")
+
     end subroutine test_math_is_close_real
 
     subroutine test_math_is_close_complex
         use stdlib_math, only: is_close
         use stdlib_error, only: check
 
-        call check(is_close((2.5,1.2), (2.5,1.2), rel_tol=1.0e-5), &
+        call check(is_close((2.5, 1.2), (2.5, 1.2), rel_tol=1.0e-5), &
                    msg="is_close((2.5,1.2), (2.5,1.2), rel_tol=1.0e-5) failed.")
-        call check(all(is_close([(2.5,1.2), (3.2,1.2)], [(2.5,1.2), (10.0,1.2)], rel_tol=1.0e-5)), &
-                   msg="all(is_close([(2.5,1.2), (3.2,1.2)], [(2.5,1.2), (10.0,1.2)], rel_tol=1.0e-5)) failed (expected).", &
+        call check(all(is_close([(2.5, 1.2), (3.2, 1.2)], [(2.5, 1.2), (10.0, 1.2)], rel_tol=1.0e-5)), &
+                   msg="all(is_close([(2.5,1.2), (3.2,1.2)], [(2.5,1.2), (10.0,1.2)], rel_tol=1.0e-5)) failed. (expected)", &
                    warn=.true.)
-        call check(all(is_close(reshape([(2.5,1.2009), (3.2,1.199999)], [1, 2]), reshape([(2.4,1.2009), (3.15,1.199999)], [1, 2]), &
-                   abs_tol=1.0e-5, rel_tol=0.1)), &
-                   msg="all(is_close(reshape([(2.5,1.2009), (3.2,1.199999)], [1, 2]), &
-                   &reshape([(2.4,1.2009), (3.15,1.199999)], [1, 2]), &
-                   &rel_tol=1.0e-5, abs_tol=0.1)) failed.")
+        call check(all(is_close(reshape([(2.5, 1.2009), (3.2, 1.199999)], [1, 2]), &
+                                reshape([(2.4, 1.2009), (3.15, 1.199999)], [1, 2]), &
+                                abs_tol=1.0e-5, rel_tol=0.1)), &
+                                msg="all(is_close(reshape([(2.5,1.2009), (3.2,1.199999)], [1, 2]), &
+                                &reshape([(2.4,1.2009), (3.15,1.199999)], [1, 2]), &
+                                &rel_tol=1.0e-5, abs_tol=0.1)) failed.")
 
         !> Tests for zeros
         call check(is_close((0.0, -0.0), (-0.0, 0.0)), msg="is_close((0.0, -0.0), (-0.0, 0.0)) failed.")
 
+        !> Tests for NaN
+        call check(is_close(cmplx(NAN, NAN), cmplx(NAN, NAN)), &
+                   msg="is_close(cmplx(NAN, NAN), cmplx(NAN, NAN)) failed. (expected)", warn=.true.)
+        call check(is_close(cmplx(NAN, NAN), cmplx(NAN, NAN), equal_nan=.true.), &
+                   msg="is_close(cmplx(NAN, NAN), cmplx(NAN, NAN), equal_nan=.true.) failed.")
+        call check(is_close(cmplx(NAN, 1.0), cmplx(NAN, 1.0)), &
+                   msg="is_close(cmplx(NAN, NAN), cmplx(NAN, NAN)) failed. (expected)", warn=.true.)
+        call check(is_close(cmplx(NAN, 1.0), cmplx(NAN, 1.0), equal_nan=.true.), &
+                   msg="is_close(cmplx(NAN, NAN), cmplx(NAN, NAN), equal_nan=.ture.) failed.")
+
     end subroutine test_math_is_close_complex
 
 end program test_math_is_close
