1
1
function LinearOperatorCollection. GradientOp (:: Type{T} ;
2
- shape:: Tuple , dim :: Union{Nothing,Int64} = nothing ) where T <: Number
3
- if dim == nothing
2
+ shape:: Tuple , dims = nothing ) where T <: Number
3
+ if dims == nothing
4
4
return GradientOpImpl (T, shape)
5
5
else
6
- return GradientOpImpl (T, shape, dim )
6
+ return GradientOpImpl (T, shape, dims )
7
7
end
8
8
end
9
9
10
-
11
- """
12
- GradientOpImpl(T::Type, shape::NTuple{1,Int64})
13
-
14
- 1d gradient operator for an array of size `shape`
15
10
"""
16
- GradientOpImpl (T:: Type , shape:: NTuple{1 ,Int64} ) = GradientOpImpl (T,shape, 1 )
11
+ GradOp (T::Type, shape::NTuple{N ,Int64})
17
12
13
+ Nd gradient operator for an array of size `shape`
18
14
"""
19
- GradientOpImpl(T::Type, shape::NTuple{2,Int64})
20
-
21
- 2d gradient operator for an array of size `shape`
22
- """
23
- function GradientOpImpl (T:: Type , shape:: NTuple{2,Int64} )
24
- return vcat ( GradientOpImpl (T,shape,1 ), GradientOpImpl (T,shape,2 ) )
15
+ function GradientOpImpl (T:: Type , shape)
16
+ shape = typeof (shape) <: Number ? (shape,) : shape # convert Number to Tuple
17
+ return vcat ([GradientOpImpl (T, shape, i) for i ∈ eachindex (shape)]. .. )
25
18
end
26
19
27
20
"""
28
- GradientOpImpl (T::Type, shape::NTuple{3 ,Int64})
21
+ GradOp (T::Type, shape::NTuple{N ,Int64}, dims )
29
22
30
- 3d gradient operator for an array of size `shape`
31
- """
32
- function GradientOpImpl (T:: Type , shape:: NTuple{3,Int64} )
33
- return vcat ( GradientOpImpl (T,shape,1 ), GradientOpImpl (T,shape,2 ), GradientOpImpl (T,shape,3 ) )
34
- end
35
-
36
- """
37
- gradOp(T::Type, shape::NTuple{N,Int64}, dim::Int64) where N
38
-
39
- directional gradient operator along the dimension `dim`
23
+ directional gradient operator along the dimensions `dims`
40
24
for an array of size `shape`
41
25
"""
42
- function GradientOpImpl (T:: Type , shape:: NTuple{N,Int64} , dim:: Int64 ) where N
26
+ function GradientOpImpl (T:: Type , shape:: NTuple{N,Int64} , dims) where N
27
+ return vcat ([GradientOpImpl (T, shape, dim) for dim ∈ dims]. .. )
28
+ end
29
+ function GradientOpImpl (T:: Type , shape:: NTuple{N,Int64} , dim:: Integer ) where N
43
30
nrow = div ( (shape[dim]- 1 )* prod (shape), shape[dim] )
44
31
ncol = prod (shape)
45
32
return LinearOperator {T} (nrow, ncol, false , false ,
46
- (res,x) -> (grad! (res,x,shape,dim) ),
47
- (res,x) -> (grad_t! (res,x,shape,dim) ),
33
+ (res,x) -> (grad! (res,x,shape,dim) ),
34
+ (res,x) -> (grad_t! (res,x,shape,dim) ),
48
35
nothing )
49
36
end
50
37
51
38
# directional gradients
52
- function grad! (res:: T , img:: U , shape:: NTuple{1,Int64} , dim:: Int64 ) where {T<: AbstractVector ,U<: AbstractVector }
53
- res .= img[1 : end - 1 ]. - img[2 : end ]
54
- end
55
-
56
- function grad! (res:: T , img:: U , shape:: NTuple{2,Int64} , dim:: Int64 ) where {T<: AbstractVector ,U<: AbstractVector }
57
- img = reshape (img,shape)
39
+ function grad! (res:: T , img:: U , shape, dim) where {T<: AbstractVector , U<: AbstractVector }
40
+ img_ = reshape (img,shape)
58
41
59
- if dim== 1
60
- res .= vec (img[1 : end - 1 ,:]. - img[2 : end ,:])
61
- else
62
- res .= vec (img[:,1 : end - 1 ]. - img[:,2 : end ])
63
- end
64
- end
42
+ δ = zeros (Int, length (shape))
43
+ δ[dim] = 1
44
+ δ = Tuple (δ)
45
+ di = CartesianIndex (δ)
65
46
66
- function grad! (res:: T ,img:: U , shape:: NTuple{3,Int64} , dim:: Int64 ) where {T<: AbstractVector ,U<: AbstractVector }
67
- img = reshape (img,shape)
47
+ res_ = reshape (res, shape .- δ)
68
48
69
- if dim== 1
70
- res .= vec (img[1 : end - 1 ,:,:]. - img[2 : end ,:,:])
71
- elseif dim== 2
72
- res.= vec (img[:,1 : end - 1 ,:]. - img[:,2 : end ,:])
73
- else
74
- res.= vec (img[:,:,1 : end - 1 ]. - img[:,:,2 : end ])
49
+ Threads. @threads for i ∈ CartesianIndices (res_)
50
+ @inbounds res_[i] = img_[i] - img_[i + di]
75
51
end
76
52
end
77
53
54
+
78
55
# adjoint of directional gradients
79
- function grad_t! (res:: T , g:: U , shape:: NTuple{1 ,Int64} , dim:: Int64 ) where {T<: AbstractVector ,U<: AbstractVector }
80
- res . = zero ( eltype (g ))
81
- res[ 1 : shape[ 1 ] - 1 ] . = g
82
- res[ 2 : shape[ 1 ]] .- = g
83
- end
56
+ function grad_t! (res:: T , g:: U , shape:: NTuple{N ,Int64} , dim:: Int64 ) where {T<: AbstractVector , U<: AbstractVector , N }
57
+ δ = zeros (Int, length (shape ))
58
+ δ[dim] = 1
59
+ δ = Tuple (δ)
60
+ di = CartesianIndex (δ)
84
61
85
- function grad_t! (res:: T , g:: U , shape:: NTuple{2,Int64} , dim:: Int64 ) where {T<: AbstractVector ,U<: AbstractVector }
86
- res .= zero (eltype (g))
87
62
res_ = reshape (res,shape)
63
+ g_ = reshape (g, shape .- δ)
88
64
89
- if dim== 1
90
- g = reshape (g,shape[1 ]- 1 ,shape[2 ])
91
- res_[1 : shape[1 ]- 1 ,:] .= g
92
- res_[2 : shape[1 ],:] .- = g
93
- else
94
- g = reshape (g,shape[1 ],shape[2 ]- 1 )
95
- res_[:,1 : shape[2 ]- 1 ] .= g
96
- res_[:,2 : shape[2 ]] .- = g
65
+ res_ .= 0
66
+ Threads. @threads for i ∈ CartesianIndices (g_)
67
+ @inbounds res_[i] = g_[i]
97
68
end
98
- end
99
-
100
- function grad_t! (res:: T , g:: U , shape:: NTuple{3,Int64} , dim:: Int64 ) where {T<: AbstractVector ,U<: AbstractVector }
101
- res .= zero (eltype (g))
102
- res_ = reshape (res,shape)
103
-
104
- if dim== 1
105
- g = reshape (g,shape[1 ]- 1 ,shape[2 ],shape[3 ])
106
- res_[1 : shape[1 ]- 1 ,:,:] .= g
107
- res_[2 : shape[1 ],:,:] .- = g
108
- elseif dim== 2
109
- g = reshape (g,shape[1 ],shape[2 ]- 1 ,shape[3 ])
110
- res_[:,1 : shape[2 ]- 1 ,:] .= g
111
- res_[:,2 : shape[2 ],:] .- = g
112
- else
113
- g = reshape (g,shape[1 ],shape[2 ],shape[3 ]- 1 )
114
- res_[:,:,1 : shape[3 ]- 1 ] .= g
115
- res_[:,:,2 : shape[3 ]] .- = g
69
+ Threads. @threads for i ∈ CartesianIndices (g_)
70
+ @inbounds res_[i + di] -= g_[i]
116
71
end
117
- end
72
+ end
0 commit comments