add missing files

Author: tknopp

src/DCTOp.jl

@@ -0,0 +1,73 @@
+export DCTOp
+
+mutable struct DCTOp{T} <: AbstractLinearOperator{T}
+  nrow :: Int
+  ncol :: Int
+  symmetric :: Bool
+  hermitian :: Bool
+  prod! :: Function
+  tprod! :: Nothing
+  ctprod! :: Function
+  nprod :: Int
+  ntprod :: Int
+  nctprod :: Int
+  args5 :: Bool
+  use_prod5! :: Bool
+  allocated5 :: Bool
+  Mv5 :: Vector{T}
+  Mtu5 :: Vector{T}
+  plan
+  dcttype::Int
+end
+
+LinearOperators.storage_type(op::DCTOp) = typeof(op.Mv5)
+
+"""
+  DCTOp(T::Type, shape::Tuple, dcttype=2)
+
+returns a `DCTOp <: AbstractLinearOperator` which performs a DCT on a given input array.
+
+# Arguments:
+* `T::Type`       - type of the array to transform
+* `shape::Tuple`  - size of the array to transform
+* `dcttype`       - type of DCT (currently `2` and `4` are supported)
+"""
+function DCTOp(T::Type, shape::Tuple, dcttype=2)
+
+  tmp=Array{Complex{real(T)}}(undef, shape) 
+  if dcttype == 2
+    plan = plan_dct!(tmp)
+    iplan = plan_idct!(tmp)
+    prod! = (res, x)  -> dct_multiply2(res, plan, x, tmp)
+    tprod! = (res, x)  -> dct_multiply2(res, iplan, x, tmp)
+
+  elseif dcttype == 4
+    factor = T(sqrt(1.0/(prod(shape)* 2^length(shape)) ))
+    plan = FFTW.plan_r2r!(tmp,FFTW.REDFT11)
+    prod! = (res, x) -> dct_multiply4(res, plan, x, tmp, factor)
+    tprod! = (res, x) -> dct_multiply4(res, plan, x, tmp, factor)
+  else
+    error("DCT type $(dcttype) not supported")
+  end
+
+  return DCTOp{T}(prod(shape), prod(shape), false, false,
+                      prod!, nothing, tprod!,
+                      0, 0, 0, true, false, true, T[], T[],
+                      plan, dcttype)
+end
+
+function dct_multiply2(res::Vector{T}, plan::P, x::Vector{T}, tmp::Array{T,D}) where {T,P,D}
+  tmp[:] .= x
+  plan * tmp
+  res .= vec(tmp)
+end
+
+function dct_multiply4(res::Vector{T}, plan::P, x::Vector{T}, tmp::Array{T,D}, factor::T) where {T,P,D}
+  tmp[:] .= x
+  plan * tmp
+  res .= factor.*vec(tmp)
+end
+
+function Base.copy(S::DCTOp)
+  return DCTOp(eltype(S), size(S.plan), S.dcttype)
+end

src/DSTOp.jl

@@ -0,0 +1,77 @@
+export DSTOp
+
+mutable struct DSTOp{T} <: AbstractLinearOperator{T}
+  nrow :: Int
+  ncol :: Int
+  symmetric :: Bool
+  hermitian :: Bool
+  prod! :: Function
+  tprod! :: Nothing
+  ctprod! :: Function
+  nprod :: Int
+  ntprod :: Int
+  nctprod :: Int
+  args5 :: Bool
+  use_prod5! :: Bool
+  allocated5 :: Bool
+  Mv5 :: Vector{T}
+  Mtu5 :: Vector{T}
+  plan
+  iplan
+end
+
+LinearOperators.storage_type(op::DSTOp) = typeof(op.Mv5)
+
+"""
+  DSTOp(T::Type, shape::Tuple)
+
+returns a `LinearOperator` which performs a DST on a given input array.
+
+# Arguments:
+* `T::Type`       - type of the array to transform
+* `shape::Tuple`  - size of the array to transform
+"""
+function DSTOp(T::Type, shape::Tuple)
+  tmp=Array{Complex{real(T)}}(undef, shape)
+
+  plan = FFTW.plan_r2r!(tmp,FFTW.RODFT10)
+  iplan = FFTW.plan_r2r!(tmp,FFTW.RODFT01)
+
+  w = weights(shape, T)
+
+  return DSTOp{T}(prod(shape), prod(shape), true, false
+            , (res,x) -> dst_multiply!(res,plan,x,tmp,w)
+            , nothing
+            , (res,x) -> dst_bmultiply!(res,iplan,x,tmp,w)
+            , 0, 0, 0, true, false, true, T[], T[]
+            , plan
+            , iplan)
+end
+
+function weights(s, T::Type)
+  w = ones(T,s...)./T(sqrt(8*prod(s)))
+  w[s[1],:,:]./= T(sqrt(2))
+  if length(s)>1
+    w[:,s[2],:]./= T(sqrt(2))
+    if length(s)>2
+      w[:,:,s[3]]./= T(sqrt(2))
+    end
+  end
+  return reshape(w,prod(s))
+end
+
+function dst_multiply!(res::Vector{T}, plan::P, x::Vector{T}, tmp::Array{T,D}, weights::Vector{T}) where {T,P,D}
+  tmp[:] .= x
+  plan * tmp
+  res .= vec(tmp).*weights
+end
+
+function dst_bmultiply!(res::Vector{T}, plan::P, x::Vector{T}, tmp::Array{T,D}, weights::Vector{T}) where {T,P,D}
+  tmp[:] .= x./weights
+  plan * tmp
+  res[:] .= vec(tmp)./(8*length(tmp))
+end
+
+function Base.copy(S::DSTOp)
+  return DSTOp(eltype(S), size(S.plan))
+end

src/FFTOp.jl

@@ -0,0 +1,96 @@
+export FFTOp
+import Base.copy
+
+mutable struct FFTOp{T} <: AbstractLinearOperator{T}
+  nrow :: Int
+  ncol :: Int
+  symmetric :: Bool
+  hermitian :: Bool
+  prod! :: Function
+  tprod! :: Nothing
+  ctprod! :: Function
+  nprod :: Int
+  ntprod :: Int
+  nctprod :: Int
+  args5 :: Bool
+  use_prod5! :: Bool
+  allocated5 :: Bool
+  Mv5 :: Vector{T}
+  Mtu5 :: Vector{T}
+  plan
+  iplan
+  shift::Bool
+  unitary::Bool
+end
+
+LinearOperators.storage_type(op::FFTOp) = typeof(op.Mv5)
+
+"""
+  FFTOp(T::Type, shape::Tuple, shift=true, unitary=true)
+
+returns an operator which performs an FFT on Arrays of type T
+
+# Arguments:
+* `T::Type`       - type of the array to transform
+* `shape::Tuple`  - size of the array to transform
+* (`shift=true`)  - if true, fftshifts are performed
+* (`unitary=true`)  - if true, FFT is normalized such that it is unitary
+"""
+function FFTOp(T::Type, shape::NTuple{D,Int64}, shift::Bool=true; unitary::Bool=true, cuda::Bool=false) where D
+  
+  #tmpVec = cuda ? CuArray{T}(undef,shape) : Array{Complex{real(T)}}(undef, shape)
+  tmpVec = Array{Complex{real(T)}}(undef, shape)
+  plan = plan_fft!(tmpVec; flags=FFTW.MEASURE)
+  iplan = plan_bfft!(tmpVec; flags=FFTW.MEASURE)
+
+  if unitary
+    facF = T(1.0/sqrt(prod(shape)))
+    facB = T(1.0/sqrt(prod(shape)))
+  else
+    facF = T(1.0)
+    facB = T(1.0)
+  end
+
+  let shape_=shape, plan_=plan, iplan_=iplan, tmpVec_=tmpVec, facF_=facF, facB_=facB
+
+  if shift
+    return FFTOp{T}(prod(shape), prod(shape), false, false
+              , (res, x) -> fft_multiply_shift!(res, plan_, x, shape_, facF_, tmpVec_) 
+              , nothing
+              , (res, x) -> fft_multiply_shift!(res, iplan_, x, shape_, facB_, tmpVec_) 
+              , 0, 0, 0, true, false, true, T[], T[]
+              , plan
+              , iplan
+              , shift
+              , unitary)
+  else
+    return FFTOp{T}(prod(shape), prod(shape), false, false
+            , (res, x) -> fft_multiply!(res, plan_, x, facF_, tmpVec_) 
+            , nothing
+            , (res, x) -> fft_multiply!(res, iplan_, x, facB_, tmpVec_)
+            , 0, 0, 0, true, false, true, T[], T[]
+            , plan
+            , iplan
+            , shift
+            , unitary)
+  end
+  end
+end
+
+function fft_multiply!(res::AbstractVector{T}, plan::P, x::AbstractVector{Tr}, factor::T, tmpVec::Array{T,D}) where {T, Tr, P<:AbstractFFTs.Plan, D}
+  tmpVec[:] .= x
+  plan * tmpVec
+  res .= factor .* vec(tmpVec)
+end
+
+function fft_multiply_shift!(res::AbstractVector{T}, plan::P, x::AbstractVector{Tr}, shape::NTuple{D}, factor::T, tmpVec::Array{T,D}) where {T, Tr, P<:AbstractFFTs.Plan, D}
+  ifftshift!(tmpVec, reshape(x,shape))
+  plan * tmpVec
+  fftshift!(reshape(res,shape), tmpVec)
+  res .*= factor
+end
+
+
+function Base.copy(S::FFTOp)
+  return FFTOp(eltype(S), size(S.plan), S.shift, unitary=S.unitary)
+end

src/GradientOp.jl

@@ -0,0 +1,109 @@
+export GradientOp
+
+"""
+    gradOp(T::Type, shape::NTuple{1,Int64})
+
+1d gradient operator for an array of size `shape`
+"""
+GradientOp(T::Type, shape::NTuple{1,Int64}) = GradientOp(T,shape,1)
+
+"""
+    gradOp(T::Type, shape::NTuple{2,Int64})
+
+2d gradient operator for an array of size `shape`
+"""
+function GradientOp(T::Type, shape::NTuple{2,Int64})
+  return vcat( GradientOp(T,shape,1), GradientOp(T,shape,2) ) 
+end
+
+"""
+    gradOp(T::Type, shape::NTuple{3,Int64})
+
+3d gradient operator for an array of size `shape`
+"""
+function GradientOp(T::Type, shape::NTuple{3,Int64})
+  return vcat( GradientOp(T,shape,1), GradientOp(T,shape,2), GradientOp(T,shape,3) ) 
+end
+
+"""
+    gradOp(T::Type, shape::NTuple{N,Int64}, dim::Int64) where N
+
+directional gradient operator along the dimension `dim`
+for an array of size `shape`
+"""
+function GradientOp(T::Type, shape::NTuple{N,Int64}, dim::Int64) 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) ), 
+                          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)
+
+  if dim==1
+    res .= vec(img[1:end-1,:].-img[2:end,:])
+  else
+    res .= vec(img[:,1:end-1].-img[:,2:end])
+  end
+end
+
+function grad!(res::T,img::U, shape::NTuple{3,Int64}, dim::Int64) where {T<:AbstractVector,U<:AbstractVector}
+  img = reshape(img,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])
+  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{2,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])
+    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
+  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
+  end
+end