stdlib/src/math/stdlib_math_diff.fypp at 99a460ac1a5351a7fb17f574edefe05c3f4af04c · srinjoy933/stdlib · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
!> 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
submodule (stdlib_math) stdlib_math_diff

    implicit none

contains

    !> `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, size_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)
        size_work = size_x + size_prepend + size_append

        if (size_work <= n_) then
            allocate(y(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
            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)
        end do

        y = work(1:size_work-n_)
        end block

    end function diff_1_${k1}$

    pure module function diff_2_${k1}$(x, n, dim, prepend, append) result(y)
        ${t1}$, intent(in) :: x(:, :)
        integer, intent(in), optional :: n, dim
        ${t1}$, intent(in), optional :: prepend(:, :), append(:, :)
        ${t1}$, allocatable :: y(:, :)
        integer :: size_prepend, size_append, size_x, size_work
        integer :: n_, dim_, i

        n_ = optval(n, 1)
        if (n_ <= 0) then
            y = x
            return
        end if

        size_prepend = 0
        size_append = 0

        dim_ = 1
        if (present(dim)) then
            if (dim == 1 .or. dim == 2) dim_ = dim
        end if

        if (present(prepend)) then
            if (size(prepend, 3 - dim_) /= size(x, 3 - dim_)) then
                error stop "stdlib_math_diff: non-differencing dimension of 'prepend' must match 'x'"
            end if
            size_prepend = size(prepend, dim_)
        end if

        if (present(append)) then
            if (size(append, 3 - dim_) /= size(x, 3 - dim_)) then
                error stop "stdlib_math_diff: non-differencing dimension of 'append' must match 'x'"
            end if
            size_append = size(append, dim_)
        end if

        size_x = size(x, dim_)
        size_work = size_x + size_prepend + size_append

        if (size_work <= n_) then
            select case (dim_)
            case (1)
                allocate(y(0, size(x, 2)))
            case (2)
                allocate(y(size(x, 1), 0))
            end select
            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
            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-i, :) = work(2:size_work-i+1, :) - work(1:size_work-i, :)
            end do

            y = work(1:size_work-n_, :)
            end block

        elseif (dim_ == 2) then
            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-i) = work(:, 2:size_work-i+1) - work(:, 1:size_work-i)
            end do

            y = work(:, 1:size_work-n_)
            end block

        end if

    end function diff_2_${k1}$
    #:endfor

end submodule stdlib_math_diff