Update implementaion of diff procedure.

zoziha · zoziha · commit ec00dfeb3ecb · 2022-01-02T09:35:19.000+08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -22,6 +22,7 @@ Features available from the latest git source
   - new procedures `arg`, `argd` and `argpi`
     [#498](https://github.com/fortran-lang/stdlib/pull/498)
   - new procedure `diff`
+    [#605](https://github.com/fortran-lang/stdlib/pull/605)
 
 Changes to existing modules
 
diff --git a/doc/specs/stdlib_math.md b/doc/specs/stdlib_math.md
@@ -654,15 +654,17 @@ end program demo_math_all_close
 
 #### Description
 
-Computes differences between adjacent elements of the input array.
+Computes differences between adjacent elements of an array.
 
 #### Syntax
 
+For rank-1 array
 ```fortran
-!> For rank-1 array
-Y = [[stdlib_math(module):diff(interface)]](X [, n])
-!> and for rank-2 array
-Y = [[stdlib_math(module):diff(interface)]](X [, n, dim])
+y = [[stdlib_math(module):diff(interface)]](x [, n, prepend, append])
+```
+and for rank-2 array
+```fortran
+y = [[stdlib_math(module):diff(interface)]](x [, n, dim, prepend, append])
 ```
 
 #### Status
@@ -675,10 +677,9 @@ Pure function.
 
 #### Arguments
 
-Note: If the value of `dim` is not equal to `1` or `2`,
-`1` will be used by the internal process of `diff`. (Not recommended)
+Note: The `x`, `prepend` and `append` arguments must have same `type`, `kind` and `rank`.
 
-`X`: Shall be a `real/integer` and `rank-1/rank-2` array.
+`x`: Shall be a `real/integer` and `rank-1/rank-2` array.
 This argument is `intent(in)`.
 
 `n`: Shall be an `integer` scalar.
@@ -689,9 +690,20 @@ It represents to calculate the n-th order difference.
 This argument is `intent(in)` and `optional`, which is `1` by default.
 It represents to calculate the difference along which dimension.
 
+`prepend`: Shall be a `real/integer` and `rank-1/rank-2` array, which is no value by default.
+This argument is `intent(in)` and `optional`.
+
+`append`: Shall be a `real/integer` and `rank-1/rank-2` array, which is no value by default.
+This argument is `intent(in)` and `optional`.
+
+Note: 
+- If the value of `n` is less than or equal to `0`, the return value of `y` is `x`. (Not recommended)
+- If the value of `dim` is not equal to `1` or `2`,
+`1` will be used by the internal process of `diff`. (Not recommended)
+
 #### Result value
 
-Note: That `Y` has one fewer element than `X`.
+Note: That `y` generally has one fewer element than `x`.
 
 Returns a `real/integer` and `rank-1/rank-2` array.
 
diff --git a/src/stdlib_math.fypp b/src/stdlib_math.fypp
@@ -362,20 +362,22 @@ module stdlib_math
     
     !> Version: experimental
     !>
-    !> Computes differences between adjacent elements of the input array.
+    !> Computes differences between adjacent elements of an array.
     !> ([Specification](../page/specs/stdlib_math.html#diff))
     interface diff
         #:set RI_KINDS_TYPES = REAL_KINDS_TYPES + INT_KINDS_TYPES
         #:for k1, t1 in RI_KINDS_TYPES
-        pure module function diff_1_${k1}$(x, n) result(Y)
+        pure module function diff_1_${k1}$(x, n, prepend, append) result(y)
             ${t1}$, intent(in) :: x(:)
             integer, intent(in), optional :: n
-            ${t1}$, allocatable :: Y(:)
+            ${t1}$, intent(in), optional :: prepend(:), append(:)
+            ${t1}$, allocatable :: y(:)
         end function diff_1_${k1}$
-        pure module function diff_2_${k1}$(X, n, dim) result(Y)
-            ${t1}$, intent(in) :: X(:, :)
+        pure module function diff_2_${k1}$(X, n, dim, prepend, append) result(y)
+            ${t1}$, intent(in) :: x(:, :)
             integer, intent(in), optional :: n, dim
-            ${t1}$, allocatable :: Y(:,:)
+            ${t1}$, intent(in), optional :: prepend(:, :), append(:, :)
+            ${t1}$, allocatable :: y(:, :)
         end function diff_2_${k1}$
         #:endfor
     end interface diff
diff --git a/src/stdlib_math_diff.fypp b/src/stdlib_math_diff.fypp
@@ -1,4 +1,4 @@
-#! Originally written in 2016 by Keurfon Luu (keurfonluu@outlook.com)
+#! Inspired by original code (MIT license) written in 2016 by Keurfon Luu (keurfonluu@outlook.com)
 #! https://github.com/keurfonluu/Forlab/blob/master/src/lib/forlab.f90#L2673
 
 #:include "common.fypp"
@@ -9,31 +9,65 @@ submodule (stdlib_math) stdlib_math_diff
 
 contains
 
-    #! `diff` computes differences of arrays of the ${t1}$ type.
+    #! `diff` computes differences of adjacent elements of an array of the ${t1}$ type.
     
     #:for k1, t1 in RI_KINDS_TYPES
-    pure module function diff_1_${k1}$(X, n) result(Y)
-        ${t1}$, intent(in) :: X(:)
+    pure module function diff_1_${k1}$(x, n, prepend, append) result(y)
+        ${t1}$, intent(in) :: x(:)
         integer, intent(in), optional :: n
-        ${t1}$, allocatable :: Y(:)
+        ${t1}$, intent(in), optional :: prepend(:), append(:)
+        ${t1}$, allocatable :: y(:)
+        integer :: size_prepend, size_append, size_x
+        ${t1}$, allocatable :: work(:)
         integer :: n_, i
 
         n_ = optval(n, 1)
+        if (n_ <= 0) then
+            y = x
+            return
+        end if
+        
+        size_prepend = 0
+        size_append = 0
+        if (present(prepend)) size_prepend = size(prepend) 
+        if (present(append)) size_append = size(append)
+        size_x = size(x)
+        
+        if (size_x + size_prepend + size_append <= n_) then
+            allocate(y(0))
+            return
+        end if
 
-        Y = X
+        allocate(work(size_x + size_prepend + size_append))
+        if (size_prepend > 0) work(:size_prepend) = prepend
+        work(size_prepend+1:size_prepend+size_x) = x
+        if (size_append > 0) work(size_prepend+size_x+1:) = append
+        
         do i = 1, n_
-            Y = Y(2:) - Y(:size(Y) - 1)
+            work(1:size(work)-i) = work(2:size(work)-i+1) - work(1:size(work)-i)
         end do
+        
+        y = work(1:size(work)-n_)
 
     end function diff_1_${k1}$
 
-    pure module function diff_2_${k1}$(X, n, dim) result(Y)
-        ${t1}$, intent(in) :: X(:, :)
+    pure module function diff_2_${k1}$(x, n, dim, prepend, append) result(y)
+        ${t1}$, intent(in) :: x(:, :)
         integer, intent(in), optional :: n, dim
-        ${t1}$, allocatable :: Y(:, :)
+        ${t1}$, intent(in), optional :: prepend(:, :), append(:, :)
+        ${t1}$, allocatable :: y(:, :)
+        integer :: size_prepend, size_append, size_x
         integer :: n_, dim_, i
+        ${t1}$, allocatable :: work(:, :)
 
         n_ = optval(n, 1)
+        if (n_ <= 0) then
+            y = x
+            return
+        end if
+        
+        size_prepend = 0
+        size_append = 0
         if (present(dim)) then
             if (dim == 1 .or. dim == 2) then
                 dim_ = dim
@@ -44,15 +78,37 @@ contains
             dim_ = 1
         end if
         
-        Y = X
+        if (present(prepend)) size_prepend = size(prepend, dim_)
+        if (present(append)) size_append = size(append, dim_)
+        size_x = size(x, dim_)
+        
+        if (size_x + size_prepend + size_append <= n_) then
+            allocate(y(0, 0))
+            return
+        end if
+        
         if (dim_ == 1) then
+            allocate(work(size_x+size_prepend+size_append, size(x, 2)))
+            if (size_prepend > 0) work(1:size_prepend, :) = prepend
+            work(size_prepend+1:size_x+size_prepend, :) = x
+            if (size_append > 0) work(size_x+size_prepend+1:, :) = append
             do i = 1, n_
-                Y = Y(2:, :) - Y(:size(Y, 1) - 1, :)
+                work(1:size(work,1)-i, :) = work(2:size(work)-i+1, :) - work(1:size(x, 1)-i, :)
             end do
-        elseif (dim == 2) then
+
+            y = work(1:size(work)-n_, :)
+            
+        elseif (dim_ == 2) then
+            allocate(work(size(x, 1), size_x+size_prepend+size_append))
+            if (size_prepend > 0) work(:, 1:size_prepend) = prepend
+            work(:, size_prepend+1:size_x+size_prepend) = x
+            if (size_append > 0) work(:, size_x+size_prepend+1:) = append
             do i = 1, n_
-                Y = Y(:, 2:) - Y(:, :size(Y, 2) - 1)
+                work(:, 1:size(work,2)-i) = work(:, 2:size(work,2)-i+1) - work(:, 1:size(work, 2)-i)
             end do
+            
+            y = work(:, 1:size(work,2)-n_)
+            
         end if
 
     end function diff_2_${k1}$
diff --git a/src/tests/math/test_stdlib_math.fypp b/src/tests/math/test_stdlib_math.fypp
@@ -378,18 +378,43 @@ contains
         type(error_type), allocatable, intent(out) :: error
         ${t1}$ :: x(6) = [${t1}$ :: 0, 5, 15, 30, 50, 75]
         ${t1}$ :: A(1, 3) = reshape([${t1}$ :: 1, 3, 5], [1, 3])
+        ${t1}$ :: B(2) = [${t1}$ :: 1, 2]
         
+        #! rank-1 diff
         call check(error, all_close(diff(x), [${t1}$ :: 5, 10, 15, 20, 25]), &
-            "diff(x) in test_diff_real_${k1}$ failed")
+            "diff(<rank-1>) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
-        
+        call check(error, all_close(diff(x, n=0), x), &
+            "diff(<rank-1>, n=0) in test_diff_real_${k1}$ failed")
         call check(error, all_close(diff(x, n=2), [${t1}$ :: 5, 5, 5, 5]), &
-            "diff(x, n=2) in test_diff_real_${k1}$ failed")
+            "diff(<rank-1>, n=2) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
         
+        call check(error, all_close(diff(x, prepend=[${t1}$ :: 1]), [${t1}$ :: -1, 5, 10, 15, 20, 25]), &
+            "diff(<rank-1>, prepend=[${t1}$ :: 1]) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
+        call check(error, all_close(diff(x, append=[${t1}$ :: 1]), [${t1}$ :: 5, 10, 15, 20, 25, -74]), &
+            "diff(<rank-1>, append=[${t1}$ :: 1]) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
+        
+        #! rank-2 diff
+        call check(error, all_close(diff(reshape(A, [3,1]), n=1, dim=1), reshape([${t1}$ :: 2, 2], [2, 1])), &
+            "diff(<rank-2>, n=1, dim=1) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
         call check(error, all_close(diff(A, n=1, dim=2), reshape([${t1}$ :: 2, 2], [1, 2])), &
-            "diff(x, n=1, dim=2) in test_diff_real_${k1}$ failed")
+            "diff(<rank-2>, n=1, dim=2) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
+        
+        call check(error, all_close(diff(A, n=1, dim=2, prepend=reshape([${t1}$ :: 1], [1, 1]), &
+            append=reshape([${t1}$ :: 2], [1, 1])), reshape([${t1}$ :: 0, 2, 2, -3], [1, 4])), &
+            "diff(<rank-2>, n=1, dim=2, prepend=reshape([${t1}$ :: 1], [1, 1]), &
+            &append=reshape([${t1}$ :: 2], [1, 1])) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
+        
+        #! size(B, dim) <= n
+        call check(error, size(diff(B, 2)), 0, "size(diff(B, 2)) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
+        call check(error, size(diff(B, 3)), 0, "size(diff(B, 3)) in test_diff_real_${k1}$ failed")
         
     end subroutine test_diff_real_${k1}$
     #:endfor
@@ -399,19 +424,39 @@ contains
         type(error_type), allocatable, intent(out) :: error
         ${t1}$ :: x(6) = [${t1}$ :: 0, 5, 15, 30, 50, 75]
         ${t1}$ :: A(1, 3) = reshape([${t1}$ :: 1, 3, 5], [1, 3])
+        ${t1}$ :: B(2) = [${t1}$ :: 1, 2]
         
+        #! rank-1 diff
         call check(error, all(diff(x) == [${t1}$ :: 5, 10, 15, 20, 25]), &
-            "diff(x) in test_diff_int_${k1}$ failed")
+            "diff(<rank-1>) in test_diff_int_${k1}$ failed")
+        if (allocated(error)) return
+        call check(error, all(diff(x, n=0) == x), &
+            "diff(<rank-1>, n=0) in test_diff_int_${k1}$ failed")
         if (allocated(error)) return
-        
         call check(error, all(diff(x, n=2) == [${t1}$ :: 5, 5, 5, 5]), &
-            "diff(x, n=2) in test_diff_int_${k1}$ failed")
+            "diff(<rank-1>, n=2) in test_diff_int_${k1}$ failed")
+        if (allocated(error)) return
+        
+        call check(error, all(diff(x, prepend=[${t1}$ :: 1]) == [${t1}$ :: -1, 5, 10, 15, 20, 25]), &
+            "diff(<rank-1>, prepend=[${t1}$ :: 1]) in test_diff_int_${k1}$ failed")
+        if (allocated(error)) return
+        call check(error, all(diff(x, append=[${t1}$ :: 1]) == [${t1}$ :: 5, 10, 15, 20, 25, -74]), &
+            "diff(<rank-1>, append=[${t1}$ :: 1]) in test_diff_int_${k1}$ failed")
         if (allocated(error)) return
         
+        #! rank-2 diff
+        call check(error, all(diff(reshape(A, [3,1]), n=1, dim=1) == reshape([${t1}$ :: 2, 2], [2, 1])), &
+            "diff(<rank-2>, n=1, dim=1) in test_diff_int_${k1}$ failed")
+        if (allocated(error)) return
         call check(error, all(diff(A, n=1, dim=2) == reshape([${t1}$ :: 2, 2], [1, 2])), &
             "diff(A, n=1, dim=2) in test_diff_int_${k1}$ failed")
         if (allocated(error)) return
         
+        #! size(B, dim) <= n
+        call check(error, size(diff(B, 2)), 0, "size(diff(B, 2)) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
+        call check(error, size(diff(B, 3)), 0, "size(diff(B, 3)) in test_diff_real_${k1}$ failed")
+        
     end subroutine test_diff_int_${k1}$
     #:endfor
 
