Improve the efficiency of diff.

Author: zoziha

doc/specs/stdlib_math.md

@@ -677,7 +677,7 @@ Pure function.
 
 #### Arguments
 
-Note: The `x`, `prepend` and `append` arguments must have same `type`, `kind` and `rank`.
+Note: The `x`, `prepend` and `append` arguments must have the same `type`, `kind` and `rank`.
 
 `x`: Shall be a `real/integer` and `rank-1/rank-2` array.
 This argument is `intent(in)`.
@@ -688,7 +688,7 @@ It represents to calculate the n-th order difference.
 
 `dim`: Shall be an `integer` scalar.
 This argument is `intent(in)` and `optional`, which is `1` by default.
-It represents to calculate the difference along which dimension.
+It gives the dimension of the input array along which the difference is calculated, between `1` and `rank(x)`.
 
 `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`.
@@ -728,5 +728,8 @@ program demo_diff
     print *, Y(2, :)        !! [4, 2]
     print *, Y(3, :)        !! [2, 4]
 
+    print *, diff(i, prepend=[0]) !! [1, 0, 1, 1, 2, 3, 5]
+    print *, diff(i, append=[21]) !! [0, 1, 1, 2, 3, 5, 8]
+    
 end program demo_diff
 ```

src/stdlib_math_diff.fypp

@@ -1,5 +1,5 @@
-#! 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
+!> Inspired by original code (MIT license) written in 2016 by Keurfon Luu (keurfonluu@outlook.com)
+!> https://github.com/keurfonluu/Forlab
 
 #:include "common.fypp"
 #:set RI_KINDS_TYPES = REAL_KINDS_TYPES + INT_KINDS_TYPES
@@ -9,16 +9,15 @@ submodule (stdlib_math) stdlib_math_diff
 
 contains
 
-    #! `diff` computes differences of adjacent elements of an array of the ${t1}$ type.
+    !> `diff` computes differences of adjacent elements of an array.
 
     #:for k1, t1 in RI_KINDS_TYPES
     pure module function diff_1_${k1}$(x, n, prepend, append) result(y)
         ${t1}$, intent(in) :: x(:)
         integer, intent(in), optional :: n
         ${t1}$, intent(in), optional :: prepend(:), append(:)
         ${t1}$, allocatable :: y(:)
-        integer :: size_prepend, size_append, size_x
-        ${t1}$, allocatable :: work(:)
+        integer :: size_prepend, size_append, size_x, size_work
         integer :: n_, i
 
         n_ = optval(n, 1)
@@ -32,22 +31,31 @@ contains
         if (present(prepend)) size_prepend = size(prepend) 
         if (present(append)) size_append = size(append)
         size_x = size(x)
+        size_work = size_x + size_prepend + size_append
 
-        if (size_x + size_prepend + size_append <= n_) then
+        if (size_work <= n_) then
             allocate(y(0))
             return
         end if
 
-        allocate(work(size_x + size_prepend + size_append))
+        !> Use a quick exit for the common case, to avoid memory allocation.
+        if (size_prepend == 0 .and. size_append == 0 .and. n_ == 1) then
+            y = x(2:) - x(1:size_x-1)
+            return
+        end if
+
+        block
+        ${t1}$ :: work(size_work)
         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_
-            work(1:size(work)-i) = work(2:size(work)-i+1) - work(1:size(work)-i)
+            work(1:size_work-i) = work(2:size_work-i+1) - work(1:size_work-i)
         end do
-        
-        y = work(1:size(work)-n_)
+
+        y = work(1:size_work-n_)
+        end block
 
     end function diff_1_${k1}$
 
@@ -56,9 +64,8 @@ contains
         integer, intent(in), optional :: n, dim
         ${t1}$, intent(in), optional :: prepend(:, :), append(:, :)
         ${t1}$, allocatable :: y(:, :)
-        integer :: size_prepend, size_append, size_x
+        integer :: size_prepend, size_append, size_x, size_work
         integer :: n_, dim_, i
-        ${t1}$, allocatable :: work(:, :)
 
         n_ = optval(n, 1)
         if (n_ <= 0) then
@@ -81,33 +88,48 @@ contains
         if (present(prepend)) size_prepend = size(prepend, dim_)
         if (present(append)) size_append = size(append, dim_)
         size_x = size(x, dim_)
+        size_work = size_x + size_prepend + size_append
 
-        if (size_x + size_prepend + size_append <= n_) then
+        if (size_work <= n_) then
             allocate(y(0, 0))
             return
         end if
 
+        !> Use a quick exit for the common case, to avoid memory allocation.
+        if (size_prepend == 0 .and. size_append == 0 .and. n_ == 1) then
+            if (dim_ == 1) then
+                y = x(2:, :) - x(1:size_x-1, :)
+            elseif (dim_ == 2) then
+                y = x(:, 2:) - x(:, 1:size_x-1)
+            end if
+            return
+        end if
+        
         if (dim_ == 1) then
-            allocate(work(size_x+size_prepend+size_append, size(x, 2)))
+            block
+            ${t1}$ :: work(size_work, 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_
-                work(1:size(work,1)-i, :) = work(2:size(work)-i+1, :) - work(1:size(x, 1)-i, :)
+                work(1:size_work-i, :) = work(2:size_work-i+1, :) - work(1:size_work-i, :)
             end do
 
-            y = work(1:size(work)-n_, :)
+            y = work(1:size_work-n_, :)
+            end block
 
         elseif (dim_ == 2) then
-            allocate(work(size(x, 1), size_x+size_prepend+size_append))
+            block
+            ${t1}$ :: work(size(x, 1), size_work)
             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_
-                work(:, 1:size(work,2)-i) = work(:, 2:size(work,2)-i+1) - work(:, 1:size(work, 2)-i)
+                work(:, 1:size_work-i) = work(:, 2:size_work-i+1) - work(:, 1:size_work-i)
             end do
 
-            y = work(:, 1:size(work,2)-n_)
+            y = work(:, 1:size_work-n_)
+            end block
 
         end if
 

src/tests/math/test_stdlib_math.fypp

@@ -380,7 +380,7 @@ contains
         ${t1}$ :: A(1, 3) = reshape([${t1}$ :: 1, 3, 5], [1, 3])
         ${t1}$ :: B(2) = [${t1}$ :: 1, 2]
 
-        #! rank-1 diff
+        !> rank-1 diff
         call check(error, all_close(diff(x), [${t1}$ :: 5, 10, 15, 20, 25]), &
             "diff(<rank-1>) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
@@ -397,7 +397,7 @@ contains
             "diff(<rank-1>, append=[${t1}$ :: 1]) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
 
-        #! rank-2 diff
+        !> 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
@@ -411,7 +411,7 @@ contains
             &append=reshape([${t1}$ :: 2], [1, 1])) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
 
-        #! size(B, dim) <= n
+        !> 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")
@@ -426,7 +426,7 @@ contains
         ${t1}$ :: A(1, 3) = reshape([${t1}$ :: 1, 3, 5], [1, 3])
         ${t1}$ :: B(2) = [${t1}$ :: 1, 2]
 
-        #! rank-1 diff
+        !> rank-1 diff
         call check(error, all(diff(x) == [${t1}$ :: 5, 10, 15, 20, 25]), &
             "diff(<rank-1>) in test_diff_int_${k1}$ failed")
         if (allocated(error)) return
@@ -444,18 +444,18 @@ contains
             "diff(<rank-1>, append=[${t1}$ :: 1]) in test_diff_int_${k1}$ failed")
         if (allocated(error)) return
 
-        #! rank-2 diff
+        !> 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")
+            "diff(<rank-2>, 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")
+        !> size(B, dim) <= n
+        call check(error, size(diff(B, 2)), 0, "size(diff(B, 2)) in test_diff_int_${k1}$ failed")
         if (allocated(error)) return
-        call check(error, size(diff(B, 3)), 0, "size(diff(B, 3)) in test_diff_real_${k1}$ failed")
+        call check(error, size(diff(B, 3)), 0, "size(diff(B, 3)) in test_diff_int_${k1}$ failed")
 
     end subroutine test_diff_int_${k1}$
     #:endfor