port more general GradientOp from RegLS to LinOpCol

Author: tknopp

ext/LinearOperatorNFFTExt/NFFTOp.jl

@@ -1,9 +1,5 @@
 import Base.adjoint
 
-function LinearOperatorCollection.createLinearOperator(::Type{Op}; kargs...) where Op <: NFFTOp{T} where T <: Number
-  return NFFTOp(T; kargs...)
-end
-
 function LinearOperatorCollection.NFFTOp(::Type{T};
     shape::Tuple, nodes::AbstractMatrix{U}, toeplitz=false, oversamplingFactor=1.25, 
    kernelSize=3, kargs...) where {U <: Number, T <: Number}

src/GradientOp.jl

@@ -1,117 +1,72 @@
 function LinearOperatorCollection.GradientOp(::Type{T};
-  shape::Tuple, dim::Union{Nothing,Int64}=nothing) where T <: Number
-  if dim == nothing
+  shape::Tuple, dims=nothing) where T <: Number
+  if dims == nothing
     return GradientOpImpl(T, shape)
   else
-    return GradientOpImpl(T, shape, dim)
+    return GradientOpImpl(T, shape, dims)
   end
 end
 
-
-"""
-    GradientOpImpl(T::Type, shape::NTuple{1,Int64})
-
-1d gradient operator for an array of size `shape`
 """
-GradientOpImpl(T::Type, shape::NTuple{1,Int64}) = GradientOpImpl(T,shape,1)
+    GradOp(T::Type, shape::NTuple{N,Int64})
 
+Nd gradient operator for an array of size `shape`
 """
-    GradientOpImpl(T::Type, shape::NTuple{2,Int64})
-
-2d gradient operator for an array of size `shape`
-"""
-function GradientOpImpl(T::Type, shape::NTuple{2,Int64})
-  return vcat( GradientOpImpl(T,shape,1), GradientOpImpl(T,shape,2) ) 
+function GradientOpImpl(T::Type, shape)
+  shape = typeof(shape) <: Number ? (shape,) : shape # convert Number to Tuple
+  return vcat([GradientOpImpl(T, shape, i) for i ∈ eachindex(shape)]...)
 end
 
 """
-    GradientOpImpl(T::Type, shape::NTuple{3,Int64})
+    GradOp(T::Type, shape::NTuple{N,Int64}, dims)
 
-3d gradient operator for an array of size `shape`
-"""
-function GradientOpImpl(T::Type, shape::NTuple{3,Int64})
-  return vcat( GradientOpImpl(T,shape,1), GradientOpImpl(T,shape,2), GradientOpImpl(T,shape,3) ) 
-end
-
-"""
-    gradOp(T::Type, shape::NTuple{N,Int64}, dim::Int64) where N
-
-directional gradient operator along the dimension `dim`
+directional gradient operator along the dimensions `dims`
 for an array of size `shape`
 """
-function GradientOpImpl(T::Type, shape::NTuple{N,Int64}, dim::Int64) where N
+function GradientOpImpl(T::Type, shape::NTuple{N,Int64}, dims) where N
+  return vcat([GradientOpImpl(T, shape, dim) for dim ∈ dims]...)
+end
+function GradientOpImpl(T::Type, shape::NTuple{N,Int64}, dim::Integer) where N
   nrow = div( (shape[dim]-1)*prod(shape), shape[dim] )
   ncol = prod(shape)
   return LinearOperator{T}(nrow, ncol, false, false,
-                          (res,x) -> (grad!(res,x,shape,dim) ), 
-                          (res,x) -> (grad_t!(res,x,shape,dim) ), 
+                          (res,x) -> (grad!(res,x,shape,dim) ),
+                          (res,x) -> (grad_t!(res,x,shape,dim) ),
                           nothing )
 end
 
 # directional gradients
-function grad!(res::T, img::U, shape::NTuple{1,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
-  res .= img[1:end-1].-img[2:end]
-end
-
-function grad!(res::T, img::U, shape::NTuple{2,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
-  img = reshape(img,shape)
+function grad!(res::T, img::U, shape, dim) where {T<:AbstractVector, U<:AbstractVector}
+  img_ = reshape(img,shape)
 
-  if dim==1
-    res .= vec(img[1:end-1,:].-img[2:end,:])
-  else
-    res .= vec(img[:,1:end-1].-img[:,2:end])
-  end
-end
+  δ = zeros(Int, length(shape))
+  δ[dim] = 1
+  δ = Tuple(δ)
+  di = CartesianIndex(δ)
 
-function grad!(res::T,img::U, shape::NTuple{3,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
-  img = reshape(img,shape)
+  res_ = reshape(res, shape .- δ)
 
-  if dim==1
-    res .= vec(img[1:end-1,:,:].-img[2:end,:,:])
-  elseif dim==2
-    res.= vec(img[:,1:end-1,:].-img[:,2:end,:])
-  else
-    res.= vec(img[:,:,1:end-1].-img[:,:,2:end])
+  Threads.@threads for i ∈ CartesianIndices(res_)
+    @inbounds res_[i] = img_[i] - img_[i + di]
   end
 end
 
+
 # adjoint of directional gradients
-function grad_t!(res::T, g::U, shape::NTuple{1,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
-  res .= zero(eltype(g))
-  res[1:shape[1]-1] .= g
-  res[2:shape[1]] .-= g
-end
+function grad_t!(res::T, g::U, shape::NTuple{N,Int64}, dim::Int64) where {T<:AbstractVector, U<:AbstractVector, N}
+  δ = zeros(Int, length(shape))
+  δ[dim] = 1
+  δ = Tuple(δ)
+  di = CartesianIndex(δ)
 
-function grad_t!(res::T, g::U, shape::NTuple{2,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
-  res .= zero(eltype(g))
   res_ = reshape(res,shape)
+  g_ = reshape(g, shape .- δ)
 
-  if dim==1
-    g = reshape(g,shape[1]-1,shape[2])
-    res_[1:shape[1]-1,:] .= g
-    res_[2:shape[1],:] .-= g
-  else
-    g = reshape(g,shape[1],shape[2]-1)
-    res_[:,1:shape[2]-1] .= g
-    res_[:,2:shape[2]] .-= g
+  res_ .= 0
+  Threads.@threads for i ∈ CartesianIndices(g_)
+    @inbounds res_[i]  = g_[i]
   end
-end
-
-function grad_t!(res::T, g::U, shape::NTuple{3,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
-  res .= zero(eltype(g))
-  res_ = reshape(res,shape)
-
-  if dim==1
-    g = reshape(g,shape[1]-1,shape[2],shape[3])
-    res_[1:shape[1]-1,:,:] .= g
-    res_[2:shape[1],:,:] .-= g
-  elseif dim==2
-    g = reshape(g,shape[1],shape[2]-1,shape[3])
-    res_[:,1:shape[2]-1,:] .= g
-    res_[:,2:shape[2],:] .-= g
-  else
-    g = reshape(g,shape[1],shape[2],shape[3]-1)
-    res_[:,:,1:shape[3]-1] .= g
-    res_[:,:,2:shape[3]] .-= g
+  Threads.@threads for i ∈ CartesianIndices(g_)
+    @inbounds res_[i + di] -= g_[i]
   end
-end
+end

test/testOperators.jl

@@ -102,7 +102,7 @@ end
 
 function testGradOp1d(N=512)
   x = rand(N)
-  G = GradientOp(eltype(x), shape=size(x))
+  G = GradientOp(eltype(x); shape=size(x))
   G0 = Bidiagonal(ones(N),-ones(N-1), :U)[1:N-1,:]
 
   y = G*x
@@ -112,12 +112,12 @@ function testGradOp1d(N=512)
   xr = transpose(G)*y
   xr0 = transpose(G0)*y0
 
-  @test norm(y - y0) / norm(y0) ≈ 0 atol=0.001
+  @test norm(xr - xr0) / norm(xr0) ≈ 0 atol=0.001
 end
 
 function testGradOp2d(N=64)
   x = repeat(1:N,1,N)
-  G = GradientOp(eltype(x), shape=size(x))
+  G = GradientOp(eltype(x); shape=size(x))
   G_1d = Bidiagonal(ones(N),-ones(N-1), :U)[1:N-1,:]
 
   y = G*vec(x)
@@ -133,6 +133,37 @@ function testGradOp2d(N=64)
   @test norm(xr - xr0) / norm(xr0) ≈ 0 atol=0.001
 end
 
+function testDirectionalGradOp(N=64)
+  x = rand(ComplexF64,N,N)
+  G1 = GradientOp(eltype(x); shape=size(x), dims=1)
+  G2 = GradientOp(eltype(x); shape=size(x), dims=2)
+  G_1d = Bidiagonal(ones(N),-ones(N-1), :U)[1:N-1,:]
+
+  y1 = G1*vec(x)
+  y2 = G2*vec(x)
+  y1_ref = zeros(ComplexF64, N-1,N)
+  y2_ref = zeros(ComplexF64, N, N-1)
+  for i=1:N
+    y1_ref[:,i] .= G_1d*x[:,i]
+    y2_ref[i,:] .= G_1d*x[i,:]
+  end
+
+  @test norm(y1-vec(y1_ref)) / norm(y1_ref) ≈ 0 atol=0.001
+  @test norm(y2-vec(y2_ref)) / norm(y2_ref) ≈ 0 atol=0.001
+  
+  x1r = transpose(G1)*y1
+  x2r = transpose(G2)*y2
+
+  x1r_ref = zeros(ComplexF64, N,N)
+  x2r_ref = zeros(ComplexF64, N,N)
+  for i=1:N
+    x1r_ref[:,i] .= transpose(G_1d)*y1_ref[:,i]
+    x2r_ref[i,:] .= transpose(G_1d)*y2_ref[i,:]
+  end
+  @test norm(x1r-vec(x1r_ref)) / norm(x1r_ref) ≈ 0 atol=0.001
+  @test norm(x2r-vec(x2r_ref)) / norm(x2r_ref) ≈ 0 atol=0.001
+end
+
 function testSampling(N=64)
   x = rand(ComplexF64,N,N)
   # index-based sampling
@@ -275,7 +306,8 @@ end
   testWeighting(512)
   @info "test GradientOp"
   testGradOp1d(512)
-  testGradOp2d(64) 
+  testGradOp2d(64)
+  testDirectionalGradOp(64) 
   @info "test SamplingOp"
   testSampling(64)
   @info "test WaveletOp"