Merge pull request #19 from JuliaImageRecon/nh/nonuniformFFTs

Adapt NFFTOp to work with NonuniformFFTs.jl

Author: nHackel

Project.toml

@@ -14,6 +14,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+NonuniformFFTs = "cd96f58b-6017-4a02-bb9e-f4d81626177f"
 NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
@@ -26,6 +27,8 @@ KernelAbstractions = "0.9"
 JLArrays = "0.2"
 AbstractNFFTs = "0.9"
 LinearOperators = "2"
+NonuniformFFTs = "0.9"
+NFFT = "0.14"
 RadonKA = "0.6"
 Wavelets = "0.9, 0.10"
 Reexport = "1.0"
@@ -36,15 +39,17 @@ 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"
 
 [targets]
-test = ["Test", "FFTW", "Wavelets", "NFFT", "JLArrays", "RadonKA"]
+test = ["Test", "FFTW", "Wavelets", "NFFT", "NonuniformFFTs", "JLArrays", "RadonKA"]
 
 [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
@@ -56,17 +56,17 @@ function NFFTOpImpl(shape::Tuple, tr::AbstractMatrix{T}; toeplitz, oversamplingF
 end
 
 function produ!(y::AbstractVector, plan::AbstractNFFTPlan, x::AbstractVector) 
-  mul!(y, plan, reshape(x,plan.N))
+  mul!(y, plan, reshape(x,size_in(plan)))
 end
 
 function ctprodu!(x::AbstractVector, plan::AbstractNFFTPlan, y::AbstractVector)
-  mul!(reshape(x, plan.N), adjoint(plan), y)
+  mul!(reshape(x, size_in(plan)), adjoint(plan), y)
 end
 
 
 function Base.copy(S::NFFTOpImpl{T, vecT, P}) where {T, vecT, P}
   plan = copy(S.plan)
-  return NFFTOpImpl{T, vecT, P}(size(plan.k,2), prod(plan.N), false, false
+  return NFFTOpImpl{T, vecT, P}(size(plan.k,2), prod(size_in(plan)), false, false
               , (res,x) -> produ!(res,plan,x)
               , nothing
               , (res,y) -> ctprodu!(res,plan,y)
@@ -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,8 +130,8 @@ 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}
-  shape = nfft.plan.N
+function LinearOperatorCollection.NFFTToeplitzNormalOp(nfft::NFFTOp{T}, W=nothing; kwargs...) where {T}
+  shape = size_in(nfft.plan)
 
   tmpVec = similar(nfft.Mv5, (2 .* shape)...)
   tmpVec .= zero(T)
@@ -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
 """

test/runtests.jl

@@ -4,6 +4,7 @@ using Random
 using LinearAlgebra
 using FFTW
 using Wavelets
+using NonuniformFFTs
 using NFFT
 using RadonKA
 using JLArrays

test/testOperators.jl

@@ -440,9 +440,15 @@ end
     @info "test WaveletOp: $arrayType"
     @testset testWavelet(64,64;arrayType)
     @testset testWavelet(64,60;arrayType)
-    @info "test NFFTOp: $arrayType"
-    arrayType == JLArray || @testset testNFFT2d(;arrayType) # JLArray does not have a NFFTPlan
-    arrayType == JLArray || @testset testNFFT3d(;arrayType) # JLArray does not have a NFFTPlan
+    for backend in [NFFT.backend(), NonuniformFFTs.backend()]
+      @info "test NFFTOp with $(string(typeof(backend))): $arrayType"
+      @testset "NFFTOp with $(string(typeof(backend)))" begin
+        with(nfft_backend => backend) do
+          arrayType == JLArray || @testset testNFFT2d(;arrayType) # JLArray does not have a NFFTPlan
+          arrayType == JLArray || @testset testNFFT3d(;arrayType) # JLArray does not have a NFFTPlan
+        end
+      end
+    end
     @info "test DiagOp: $arrayType"
     @testset testDiagOp(;arrayType)
     @info "test RadonOp: $arrayType"