Add NonuniformFFTs package extension for toeplitz operator

Author: nHackel

Project.toml

@@ -39,6 +39,7 @@ AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
+NonuniformFFTs = "cd96f58b-6017-4a02-bb9e-f4d81626177f"
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 RadonKA = "86de8297-835b-47df-b249-c04e8db91db5"
@@ -48,6 +49,7 @@ test = ["Test", "FFTW", "Wavelets", "NFFT", "NonuniformFFTs", "JLArrays", "Radon
 
 [extensions]
 LinearOperatorNFFTExt = ["AbstractNFFTs", "FFTW"]
+LinearOperatorNonuniformFFTsExt = ["AbstractNFFTs", "NonuniformFFTs", "FFTW"]
 LinearOperatorFFTWExt = "FFTW"
 LinearOperatorWaveletExt = "Wavelets"
 LinearOperatorGPUArraysExt = "GPUArrays"

ext/LinearOperatorNFFTExt/NFFTOp.jl

@@ -19,7 +19,7 @@ function LinearOperatorCollection.NFFTOp(::Type{T};
   return NFFTOpImpl(shape, nodes; toeplitz, oversamplingFactor, kernelSize, kargs... )
 end
 
-mutable struct NFFTOpImpl{T, vecT, P <: AbstractNFFTPlan} <: NFFTOp{T}
+mutable struct NFFTOpImpl{T, vecT, P} <: NFFTOp{T, P}
   nrow :: Int
   ncol :: Int
   symmetric :: Bool
@@ -107,7 +107,7 @@ end
 
 LinearOperators.storage_type(op::NFFTToeplitzNormalOp) = typeof(op.Mv5)
 
-function NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1::matT, xL2::matT) where {T, D, matT <: AbstractArray{T, D}}
+function LinearOperatorCollection.NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1::matT, xL2::matT) where {T, D, matT <: AbstractArray{T, D}}
 
   function produ!(y, shape, fftplan, ifftplan, λ, xL1, xL2, x)
     xL1 .= 0
@@ -130,7 +130,7 @@ function NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1::matT, xL2::m
          , shape, W, fftplan, ifftplan, λ, xL1, xL2)
 end
 
-function NFFTToeplitzNormalOp(nfft::NFFTOp{T}, W=nothing; kwargs...) where {T}
+function LinearOperatorCollection.NFFTToeplitzNormalOp(nfft::NFFTOp{T}, W=nothing; kwargs...) where {T}
   shape = size_in(nfft.plan)
 
   tmpVec = similar(nfft.Mv5, (2 .* shape)...)
@@ -161,12 +161,12 @@ function NFFTToeplitzNormalOp(nfft::NFFTOp{T}, W=nothing; kwargs...) where {T}
   xL1 = tmpVec
   xL2 = similar(xL1)
 
-  return NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1, xL2)
+  return LinearOperatorCollection.NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1, xL2)
 end
 
 function LinearOperatorCollection.normalOperator(S::NFFTOpImpl{T}, W = nothing; copyOpsFn = copy, kwargs...) where T
   if S.toeplitz
-    return NFFTToeplitzNormalOp(S,W; kwargs...)
+    return LinearOperatorCollection.NFFTToeplitzNormalOp(S,W; kwargs...)
   else
     return NormalOp(eltype(S); parent = S, weights = W)
   end
@@ -175,5 +175,5 @@ 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))
+  return LinearOperatorCollection.NFFTToeplitzNormalOp(A.shape, A.weights, fftplan, ifftplan, A.λ, copy(A.xL1), copy(A.xL2))
 end

ext/LinearOperatorNonuniformFFTsExt/LinearOperatorNonuniformFFTsExt.jl

@@ -0,0 +1,46 @@
+module LinearOperatorNonuniformFFTsExt
+
+using LinearOperatorCollection, AbstractNFFTs, NonuniformFFTs, NonuniformFFTs.Kernels, FFTW
+
+function LinearOperatorCollection.NFFTToeplitzNormalOp(nfft::NFFTOp{T, P}, W=nothing; kwargs...) where {T, vecT, P <: NonuniformFFTs.NFFTPlan}
+  shape = size_in(nfft.plan)
+
+  tmpVec = similar(nfft.Mv5, (2 .* shape)...)
+  tmpVec .= zero(T)
+
+  # plan the FFTs
+  fftplan  = plan_fft(tmpVec; kwargs...)
+  ifftplan = plan_ifft(tmpVec; kwargs...)
+
+  # TODO extend the following function by weights
+  # λ = calculateToeplitzKernel(shape, nfft.plan.k; m = nfft.plan.params.m, σ = nfft.plan.params.σ, window = nfft.plan.params.window, LUTSize = nfft.plan.params.LUTSize, fftplan = fftplan)
+
+  shape_os = 2 .* shape
+  baseArrayType = Base.typename(typeof(tmpVec)).wrapper # https://github.com/JuliaLang/julia/issues/35543
+
+  nufft = nfft.plan.p
+  σ = nufft.σ
+  m = Kernels.half_support(first(nufft.kernels))
+  k = stack(nufft.points, dims = 1)
+
+  p = plan_nfft(baseArrayType, k, shape_os; m = m, σ = σ,
+		precompute=AbstractNFFTs.POLYNOMIAL, fftflags=FFTW.ESTIMATE, blocking=true)
+  tmpOnes = similar(tmpVec, size(k, 2))
+  tmpOnes .= one(T)
+  
+  if !isnothing(W)
+    eigMat = adjoint(p) * ( W  * tmpOnes)
+  else
+    eigMat = adjoint(p) * (tmpOnes)
+  end
+
+  λ = fftplan * fftshift(eigMat)
+
+  xL1 = tmpVec
+  xL2 = similar(xL1)
+
+  return LinearOperatorCollection.NFFTToeplitzNormalOp(shape, W, fftplan, ifftplan, λ, xL1, xL2)
+end
+
+
+end

src/LinearOperatorCollection.jl

@@ -30,20 +30,22 @@ end
 
 export linearOperatorList, createLinearOperator
 export AbstractLinearOperatorFromCollection, WaveletOp, FFTOp, DCTOp, DSTOp, NFFTOp,
-       SamplingOp, NormalOp, WeightingOp, GradientOp, RadonOp
+       SamplingOp, NormalOp, WeightingOp, GradientOp, RadonOp, NFFTToeplitzNormalOp
 
 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 NFFTOp{T, P} <: AbstractLinearOperatorFromCollection{T} end
 abstract type SamplingOp{T} <: AbstractLinearOperatorFromCollection{T} end
 abstract type NormalOp{T,S} <: AbstractLinearOperatorFromCollection{T} end
 abstract type GradientOp{T} <: AbstractLinearOperatorFromCollection{T} end
 abstract type RadonOp{T} <: AbstractLinearOperatorFromCollection{T} end
 
 
+function NFFTToeplitzNormalOp end
+
 """
   returns a list of currently implemented `LinearOperator`s
 """