port remaining Ops to new interface and prepare for NFFTOp

Author: tknopp

Project.toml

@@ -13,20 +13,23 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 [weakdeps]
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 
 [extensions]
 LinearOperatorWaveletExt = "Wavelets"
 LinearOperatorFFTWExt = "FFTW"
+LinearOperatorNFFTExt = "NFFT"
 
 [compat]
 julia = "1.6"
 FFTW = "1.0"
 LinearOperators = "2.3.3"
 Reexport = "1.0"
 Wavelets = "0.9"
+NFFT = "0.13"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "FFTW", "Wavelets"]
+test = ["Test", "FFTW", "Wavelets", "NFFT"]

ext/LinearOperatorFFTWExt/DCTOp.jl

@@ -5,7 +5,7 @@ function LinearOperatorCollection.constructLinearOperator(::Type{Op};
   return DCTOpImpl(T, shape, dcttype)
 end
 
-mutable struct DCTOpImpl{T} <: AbstractLinearOperatorFromCollection{T}
+mutable struct DCTOpImpl{T} <: DCTOp{T}
   nrow :: Int
   ncol :: Int
   symmetric :: Bool

ext/LinearOperatorFFTWExt/DSTOp.jl

@@ -5,7 +5,7 @@ function LinearOperatorCollection.constructLinearOperator(::Type{Op};
   return DSTOpImpl(T, shape)
 end
 
-mutable struct DSTOpImpl{T} <: AbstractLinearOperatorFromCollection{T}
+mutable struct DSTOpImpl{T} <: DSTOp{T}
   nrow :: Int
   ncol :: Int
   symmetric :: Bool

ext/LinearOperatorFFTWExt/FFTOp.jl

@@ -6,7 +6,7 @@ function LinearOperatorCollection.constructLinearOperator(::Type{Op};
   return FFTOpImpl(T, shape, shift; unitary, cuda)
 end
 
-mutable struct FFTOpImpl{T} <: AbstractLinearOperatorFromCollection{T}
+mutable struct FFTOpImpl{T} <: FFTOp{T}
   nrow :: Int
   ncol :: Int
   symmetric :: Bool

ext/LinearOperatorNFFTExt/LinearOperatorNFFTExt.jl

@@ -0,0 +1,7 @@
+module LinearOperatorNFFTExt
+
+using LinearOperatorCollection, NFFT
+
+include("NFFTOp.jl")
+
+end

ext/LinearOperatorNFFTExt/NFFTOp.jl

@@ -0,0 +1,166 @@
+export NFFTOpImpl
+import Base.adjoint
+
+function LinearOperatorCollection.constructLinearOperator(::Type{Op};
+    shape::Tuple, nodes::AbstractMatrix{T}, toeplitz=false, oversamplingFactor=1.25, 
+   kernelSize=3, kargs...) where Op <: NFFTOp{T} where T <: Number
+  return NFFTOpImpl(T, shape, nodes; toeplitz, oversamplingFactor, kernelSize, kargs... )
+end
+
+mutable struct NFFTOpImpl{T} <: NFFTOp{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
+  toeplitz :: Bool
+end
+
+LinearOperators.storage_type(op::NFFTOpImpl) = typeof(op.Mv5)
+
+"""
+    NFFTOpImpl(shape::Tuple, tr::Trajectory; kargs...)
+    NFFTOpImpl(shape::Tuple, tr::AbstractMatrix; kargs...)
+
+generates a `NFFTOpImpl` which evaluates the MRI Fourier signal encoding operator using the NFFT.
+
+# Arguments:
+* `shape::NTuple{D,Int64}`  - size of image to encode/reconstruct
+* `tr`                      - Either a `Trajectory` object, or a `ND x Nsamples` matrix for an ND-dimenensional (e.g. 2D or 3D) NFFT with `Nsamples` k-space samples
+* (`nodes=nothing`)         - Array containg the trajectory nodes (redundant)
+* (`kargs`)                 - additional keyword arguments
+"""
+function NFFTOpImpl(shape::Tuple, tr::AbstractMatrix{T}; toeplitz=false, oversamplingFactor=1.25, kernelSize=3, kargs...) where {T}
+
+  plan = plan_nfft(tr, shape, m=kernelSize, σ=oversamplingFactor, precompute=NFFT.TENSOR,
+		                          fftflags=FFTW.ESTIMATE, blocking=true)
+
+  return NFFTOpImpl{Complex{T}}(size(tr,2), prod(shape), false, false
+            , (res,x) -> produ!(res,plan,x)
+            , nothing
+            , (res,y) -> ctprodu!(res,plan,y)
+            , 0, 0, 0, false, false, false, Complex{T}[], Complex{T}[]
+            , plan, toeplitz)
+end
+
+function produ!(y::AbstractVector, plan::NFFT.NFFTPlan, x::AbstractVector) 
+  mul!(y, plan, reshape(x,plan.N))
+end
+
+function ctprodu!(x::AbstractVector, plan::NFFT.NFFTPlan, y::AbstractVector)
+  mul!(reshape(x, plan.N), adjoint(plan), y)
+end
+
+
+function Base.copy(S::NFFTOpImpl{T}) where {T}
+  plan = copy(S.plan)
+  return NFFTOpImpl{T}(size(plan.k,2), prod(plan.N), false, false
+              , (res,x) -> produ!(res,plan,x)
+              , nothing
+              , (res,y) -> ctprodu!(res,plan,y)
+              , 0, 0, 0, false, false, false, T[], T[]
+              , plan, S.toeplitz)
+end
+
+
+
+#########################################################################
+### Toeplitz Operator ###
+#########################################################################
+
+mutable struct NFFTToeplitzNormalOp{T,D,W} <: AbstractLinearOperator{T}
+  nrow :: Int
+  ncol :: Int
+  symmetric :: Bool
+  hermitian :: Bool
+  prod! :: Function
+  tprod! :: Nothing
+  ctprod! :: Nothing
+  nprod :: Int
+  ntprod :: Int
+  nctprod :: Int
+  args5 :: Bool
+  use_prod5! :: Bool
+  allocated5 :: Bool
+  Mv5 :: Vector{T}
+  Mtu5 :: Vector{T}
+  shape::NTuple{D,Int}
+  weights::W
+  fftplan
+  ifftplan
+  λ::Array{T}
+  xL1::Array{T,D}
+  xL2::Array{T,D}
+end
+
+LinearOperators.storage_type(op::NFFTToeplitzNormalOp) = typeof(op.Mv5)
+
+function NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1::Array{T,D}, xL2::Array{T,D}) where {T,D}
+
+  function produ!(y, shape, fftplan, ifftplan, λ, xL1, xL2, x)
+    xL1 .= 0
+    x = reshape(x, shape)
+  
+    xL1[CartesianIndices(x)] .= x
+    mul!(xL2, fftplan, xL1)
+    xL2 .*= λ
+    mul!(xL1, ifftplan, xL2)
+  
+    y .= vec(xL1[CartesianIndices(x)])
+    return y
+  end
+
+  return NFFTToeplitzNormalOp(prod(shape), prod(shape), false, false
+         , (res,x) -> produ!(res, shape, fftplan, ifftplan, λ, xL1, xL2, x)
+         , nothing
+         , nothing
+         , 0, 0, 0, false, false, false, T[], T[]
+         , shape, W, fftplan, ifftplan, λ, xL1, xL2)
+end
+
+function NFFTToeplitzNormalOp(S::NFFTOp{T}, W=opEye(T,size(S,1))) where {T}
+  shape = S.plan.N
+
+  # plan the FFTs
+  fftplan  = plan_fft( zeros(T, 2 .* shape);flags=FFTW.MEASURE)
+  ifftplan = plan_ifft(zeros(T, 2 .* shape);flags=FFTW.MEASURE)
+
+  # TODO extend the following function by weights
+  # λ = calculateToeplitzKernel(shape, S.plan.k; m = S.plan.params.m, σ = S.plan.params.σ, window = S.plan.params.window, LUTSize = S.plan.params.LUTSize, fftplan = fftplan)
+
+  shape_os = 2 .* shape
+  p = plan_nfft(typeof(S.plan.k), S.plan.k, shape_os; m = S.plan.params.m, σ = S.plan.params.σ,
+		precompute=NFFT.POLYNOMIAL, fftflags=FFTW.ESTIMATE, blocking=true)
+  eigMat = adjoint(p) * ( W  * ones(T, size(S.plan.k,2)))
+  λ = fftplan * fftshift(eigMat)
+
+  xL1 = Array{T}(undef, 2 .* shape)
+  xL2 = similar(xL1)
+
+  return NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1, xL2)
+end
+
+function LinearOperatorCollection.normalOperator(S::NFFTOpImpl{T}, W=opEye(T,size(S,1))) where T
+  if S.toeplitz
+    return NFFTToeplitzNormalOp(S,W)
+  else
+    return NormalOp(S,W)
+  end
+end
+
+function Base.copy(A::NFFTToeplitzNormalOp{T,D,W}) where {T,D,W}
+  fftplan  = plan_fft( zeros(T, 2 .* A.shape); flags=FFTW.MEASURE)
+  ifftplan = plan_ifft(zeros(T, 2 .* A.shape); flags=FFTW.MEASURE)
+  return NFFTToeplitzNormalOp(A.shape, A.weights, fftplan, ifftplan, A.λ, copy(A.xL1), copy(A.xL2))
+end

src/GradientOp.jl

@@ -1,28 +1,36 @@
-export GradientOp
+function LinearOperatorCollection.constructLinearOperator(::Type{Op};
+  shape::Tuple, dim::Union{Nothing,Int64}=nothing) where Op <: GradientOp{T} where T <: Number
+  if dim == nothing
+    return GradientOpImpl(T, shape)
+  else
+    return GradientOpImpl(T, shape, dim)
+  end
+end
+
 
 """
-    gradOp(T::Type, shape::NTuple{1,Int64})
+    GradientOpImpl(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)
+GradientOpImpl(T::Type, shape::NTuple{1,Int64}) = GradientOpImpl(T,shape,1)
 
 """
-    gradOp(T::Type, shape::NTuple{2,Int64})
+    GradientOpImpl(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) ) 
+function GradientOpImpl(T::Type, shape::NTuple{2,Int64})
+  return vcat( GradientOpImpl(T,shape,1), GradientOpImpl(T,shape,2) ) 
 end
 
 """
-    gradOp(T::Type, shape::NTuple{3,Int64})
+    GradientOpImpl(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) ) 
+function GradientOpImpl(T::Type, shape::NTuple{3,Int64})
+  return vcat( GradientOpImpl(T,shape,1), GradientOpImpl(T,shape,2), GradientOpImpl(T,shape,3) ) 
 end
 
 """
@@ -31,7 +39,7 @@ end
 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
+function GradientOpImpl(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,

src/LinearOperatorCollection.jl

@@ -6,7 +6,6 @@ import LinearAlgebra.BLAS: gemv, gemv!
 import LinearAlgebra: BlasFloat, normalize!, norm, rmul!, lmul!
 using SparseArrays
 using Random
-#using CUDA
 
 using Reexport
 @reexport using Reexport
@@ -15,10 +14,6 @@ using Reexport
 LinearOperators.use_prod5!(op::opEye) = false
 LinearOperators.has_args5(op::opEye) = false
 
-
-const Trafo = Union{AbstractMatrix, AbstractLinearOperator, Nothing}
-const FuncOrNothing = Union{Function, Nothing}
-
 # Helper function to wrap a prod into a 5-args mul
 function wrapProd(prod::Function)
   λ = (res, x, α, β) -> begin
@@ -31,68 +26,35 @@ function wrapProd(prod::Function)
   return λ
 end
 
-include("GradientOp.jl")
-include("SamplingOp.jl")
-include("WeightingOp.jl")
-include("NormalOp.jl")
-
-export linearOperator, linearOperatorList
-
-export constructLinearOperator
-export AbstractLinearOperatorFromCollection, WaveletOp, FFTOp, DCTOp, DSTOp
+export linearOperatorList, constructLinearOperator
+export AbstractLinearOperatorFromCollection, WaveletOp, FFTOp, DCTOp, DSTOp, NFFTOp,
+       SamplingOp, NormalOp, WeightingOp, GradientOp
 
 abstract type AbstractLinearOperatorFromCollection{T} <: AbstractLinearOperator{T} end
 abstract type WaveletOp{T} <: AbstractLinearOperatorFromCollection{T} end
 abstract type FFTOp{T} <: AbstractLinearOperatorFromCollection{T} end
 abstract type DCTOp{T} <: AbstractLinearOperatorFromCollection{T} end
 abstract type DSTOp{T} <: AbstractLinearOperatorFromCollection{T} end
+abstract type NFFTOp{T} <: AbstractLinearOperatorFromCollection{T} end
+abstract type SamplingOp{T} <: AbstractLinearOperatorFromCollection{T} end
+abstract type NormalOp{T} <: AbstractLinearOperatorFromCollection{T} end
+abstract type WeightingOp{T} <: AbstractLinearOperatorFromCollection{T} end
+abstract type GradientOp{T} <: AbstractLinearOperatorFromCollection{T} end
 
 function constructLinearOperator(::Type{<:AbstractLinearOperatorFromCollection}, args...; kargs...) 
   error("Operator can't be constructed. You need to load another package!")
 end
 
-linearOperator(op::Nothing, shape, T::Type=ComplexF32) = nothing
-
 """
   returns a list of currently implemented `LinearOperator`s
 """
 function linearOperatorList()
-  return ["DCT-II", "DCT-IV", "FFT", "DST", "Wavelet", "Gradient"]
+  return subtypes(AbstractLinearOperatorFromCollection)
 end
 
-"""
-    linearOperator(op::AbstractString, shape)
-
-returns the `LinearOperator` with name `op`.
-
-# valid names
-* `"FFT"`
-* `"DCT-II"`
-* `"DCT-IV"`
-* `"DST"`
-* `"Wavelet"`
-* `"Gradient"`
-"""
-function linearOperator(op::AbstractString, shape, T::Type=ComplexF32)
-  shape_ = tuple(shape...)
-  if op == "FFT"
-    trafo = FFTOp(T, shape_, false) #FFTOperator(shape)
-  elseif op == "DCT-II"
-    shape_ = tuple(shape[shape .!= 1]...)
-    trafo = DCTOp(T, shape_, 2)
-  elseif op == "DCT-IV"
-    shape_ = tuple(shape[shape .!= 1]...)
-    trafo = DCTOp(T, shape_, 4)
-  elseif op == "DST"
-    trafo = DSTOp(T, shape_)
-  elseif op == "Wavelet"
-    trafo = WaveletOp(T,shape_)
-  elseif op=="Gradient"
-    trafo = GradientOp(T,shape_)
-  else
-    error("Unknown transformation")
-  end
-  trafo
-end
+include("GradientOp.jl")
+include("SamplingOp.jl")
+include("WeightingOp.jl")
+include("NormalOp.jl")
 
 end