Backend based AbstractNFFT implementation

.github/workflows/ci.yml

@@ -47,7 +47,6 @@ jobs:
           using Pkg
           Pkg.add(PackageSpec(path=pwd(), subdir="AbstractNFFTs"))
           Pkg.add(PackageSpec(path=pwd(), subdir="NFFTTools"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="CuNFFT"))
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:
@@ -70,7 +69,6 @@ jobs:
           using Pkg
           Pkg.add(PackageSpec(path=pwd(), subdir="AbstractNFFTs"))
           Pkg.add(PackageSpec(path=pwd(), subdir="NFFTTools"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="CuNFFT"))
       - uses: julia-actions/julia-buildpkg@latest
       - uses: julia-actions/julia-docdeploy@latest
         env:

AbstractNFFTs/Project.toml

@@ -1,12 +1,13 @@
 name = "AbstractNFFTs"
 uuid = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
 author = ["Tobias Knopp <[email protected]>"]
-version = "0.8.3"
+version = "0.9.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -16,4 +17,5 @@ AbstractNFFTsChainRulesCoreExt = "ChainRulesCore"
 
 [compat]
 julia = "1.6"
-ChainRulesCore = "1"
+ChainRulesCore = "1"
+ScopedValues = "1"

AbstractNFFTs/src/AbstractNFFTs.jl

@@ -3,7 +3,16 @@ module AbstractNFFTs
 using LinearAlgebra
 using Printf
 
+# Remove this difference once 1.11 or higher becomes lower bound
+if VERSION >= v"1.11"
+  using Base.ScopedValues
+else
+  using ScopedValues
+end
+
+
 # interface
+export AbstractNFFTBackend, nfft_backend, with
 export AbstractFTPlan, AbstractRealFTPlan, AbstractComplexFTPlan,
        AbstractNFFTPlan, AbstractNFCTPlan, AbstractNFSTPlan, AbstractNNFFTPlan, 
        plan_nfft, plan_nfct, plan_nfst, mul!, size_in, size_out, nodes!
@@ -25,6 +34,7 @@ include("misc.jl")
 include("interface.jl")
 include("derived.jl")
 
+
 @static if !isdefined(Base, :get_extension)
   import Requires
 end

AbstractNFFTs/src/derived.jl

@@ -10,64 +10,189 @@ planfunc = Symbol("plan_"*"$op")
 
 # The following automatically call the plan_* version for type Array
 
-$(planfunc)(k::AbstractArray, N::Union{Integer,NTuple{D,Int}}, args...; kargs...) where {D} =
-    $(planfunc)(Array, k, N, args...; kargs...)
+$(planfunc)(b::AbstractNFFTBackend, k::AbstractArray, N::Union{Integer,NTuple{D,Int}}, args...; kargs...) where {D} =
+    $(planfunc)(b, Array, k, N, args...; kargs...)
 
-$(planfunc)(k::AbstractArray, y::AbstractArray, args...; kargs...) =
-    $(planfunc)(Array, k, y, args...; kargs...)
+$(planfunc)(b::AbstractNFFTBackend, k::AbstractArray, y::AbstractArray, args...; kargs...) =
+    $(planfunc)(b, Array, k, y, args...; kargs...)
+
+$(planfunc)(k::AbstractArray, args...; kargs...) = $(planfunc)(active_backend(), k, args...; kargs...)
 
 # The follow convert 1D parameters into the format required by the plan
 
-$(planfunc)(Q::Type, k::AbstractVector, N::Integer, rest...; kwargs...)  =
-    $(planfunc)(Q, collect(reshape(k,1,length(k))), (N,), rest...; kwargs...)
+$(planfunc)(b::AbstractNFFTBackend, Q::Type, k::AbstractVector, N::Integer, rest...; kwargs...)  =
+    $(planfunc)(b, Q, collect(reshape(k,1,length(k))), (N,), rest...; kwargs...)
+
+$(planfunc)(b::AbstractNFFTBackend, Q::Type, k::AbstractVector, N::NTuple{D,Int}, rest...; kwargs...) where {D} =
+    $(planfunc)(b, Q, collect(reshape(k,1,length(k))), N, rest...; kwargs...) 
 
-$(planfunc)(Q::Type, k::AbstractVector, N::NTuple{D,Int}, rest...; kwargs...) where {D} =
-    $(planfunc)(Q, collect(reshape(k,1,length(k))), N, rest...; kwargs...) 
+$(planfunc)(b::AbstractNFFTBackend, Q::Type, k::AbstractMatrix, N::NTuple{D,Int}, rest...; kwargs...) where {D}  =
+    $(planfunc)(b, Q, collect(k), N, rest...; kwargs...)
 
-$(planfunc)(Q::Type, k::AbstractMatrix, N::NTuple{D,Int}, rest...; kwargs...) where {D}  =
-    $(planfunc)(Q, collect(k), N, rest...; kwargs...)
+$(planfunc)(Q::Type, args...; kwargs...) = $(planfunc)(active_backend(), Q, args...; kwargs...)
 
+$(planfunc)(::Missing, args...; kwargs...) = no_backend_error()
 end
 end
 
 ## NNFFT constructor
-plan_nnfft(Q::Type, k::AbstractVector, y::AbstractVector, rest...; kwargs...)  =
-    plan_nnfft(Q, collect(reshape(k,1,length(k))), collect(reshape(y,1,length(k))), rest...; kwargs...)
-
+plan_nnfft(Q::Type, args...; kwargs...) = plan_nnfft(active_backend(), Q, args...; kwargs...)
+plan_nnfft(b::AbstractNFFTBackend, Q::Type, k::AbstractVector, y::AbstractVector, rest...; kwargs...)  =
+    plan_nnfft(b, Q, collect(reshape(k,1,length(k))), collect(reshape(y,1,length(k))), rest...; kwargs...)
+plan_nnfft(::Missing, args...; kwargs...) = no_backend_error()
 
 
 ###############################################
 # Allocating trafo functions with plan creation
 ###############################################
 
+"""
+    nfft(k, f, rest...; kwargs...)
+    nfft(backend, k, f, rest...; kwargs...)
+
+calculates the nfft of the array `f` for the nodes contained in the matrix `k`
+The output is a vector of length M=`size(nodes,2)`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+Backends can also be set with a scoped value overriding the current active backend within a scope:
+
+```julia
+julia> NFFT.activate!()
+
+julia> nfft(k, f, rest...; kwargs...) # uses NFFT
+
+julia> with(nfft_backend => NonuniformFFTs.backend()) do
+          nfft(k, f, rest...; kwargs...) # uses NonuniformFFTs
+       end
+```
+"""
+nfft
+"""
+    nfft_adjoint(k, N, fHat, rest...; kwargs...)
+    nfft_adjoint(backend, k, N, fHat, rest...; kwargs...)
+
+calculates the adjoint nfft of the vector `fHat` for the nodes contained in the matrix `k`.
+The output is an array of size `N`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+Backends can also be set with a scoped value overriding the current active backend within a scope:
+
+```julia
+julia> NFFT.activate!()
+
+julia> nfft_adjoint(k, N, fHat, rest...; kwargs...) # uses NFFT
+
+julia> with(nfft_backend => NonuniformFFTs.backend()) do
+          nfft_adjoint(k, N, fHat, rest...; kwargs...) # uses NonuniformFFTs
+       end
+```
+"""
+nfft_adjoint
+"""
+    nfft_transpose(k, N, fHat, rest...; kwargs...)
+    nfft_transpose(backend, k, N, fHat, rest...; kwargs...)
+
+calculates the transpose nfft of the vector `fHat` for the nodes contained in the matrix `k`.
+The output is an array of size `N`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+Backends can also be set with a scoped value overriding the current active backend within a scope:
+
+```julia
+julia> NFFT.activate!()
+
+julia> nfft_transpose(k, N, fHat, rest...; kwargs...) # uses NFFT
+
+julia> with(nfft_backend => NonuniformFFTs.backend()) do
+          nfft_transpose(k, N, fHat, rest...; kwargs...) # uses NonuniformFFTs
+       end
+```
+"""
+nfft_transpose
+
+"""
+    nfct(k, f, rest...; kwargs...)
+    nfct(backend, k, f, rest...; kwargs...)
+
+calculates the nfct of the array `f` for the nodes contained in the matrix `k`
+The output is a vector of length M=`size(nodes,2)`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+"""
+nfct
+"""
+    nfct_adjoint(k, N, fHat, rest...; kwargs...)
+    nfct_adjoint(backend, k, N, fHat, rest...; kwargs...)
+
+calculates the adjoint nfct of the vector `fHat` for the nodes contained in the matrix `k`.
+The output is an array of size `N`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+"""
+nfct_adjoint
+"""
+    nfct_transpose(k, N, fHat, rest...; kwargs...)
+    nfct_transpose(backend, k, N, fHat, rest...; kwargs...)
+
+calculates the transpose nfct of the vector `fHat` for the nodes contained in the matrix `k`.
+The output is an array of size `N`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+"""
+nfct_transpose
+
+"""
+    nfst(k, f, rest...; kwargs...)
+    nfst(backend, k, f, rest...; kwargs...)
+
+calculates the nfst of the array `f` for the nodes contained in the matrix `k`
+The output is a vector of length M=`size(nodes,2)`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+"""
+nfst
+"""
+    nfst_adjoint(k, N, fHat, rest...; kwargs...)
+    nfst_adjoint(backend, k, N, fHat, rest...; kwargs...)
+
+calculates the adjoint nfst of the vector `fHat` for the nodes contained in the matrix `k`.
+The output is an array of size `N`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+"""
+nfst_adjoint
+"""
+    nfst_transpose(k, N, fHat, rest...; kwargs...)
+    nfst_transpose(backend, k, N, fHat, rest...; kwargs...)
+
+calculates the transpose nfst of the vector `fHat` for the nodes contained in the matrix `k`.
+The output is an array of size `N`.
+
+Uses the active AbstractNFFTs `backend` if no `backend` argument is provided. Backends can be activated with `BackendModule.activate!()`.
+"""
+nfst_transpose
+
 for (op,trans) in zip([:nfft, :nfct, :nfst],
                       [:adjoint, :transpose, :transpose])
 planfunc = Symbol("plan_$(op)")
 tfunc = Symbol("$(op)_$(trans)")
 @eval begin 
 
-# TODO fix comments (how?)
-"""
-nfft(k, f, rest...; kwargs...)
-
-calculates the nfft of the array `f` for the nodes contained in the matrix `k`
-The output is a vector of length M=`size(nodes,2)`
-"""
-function $(op)(k, f::AbstractArray; kargs...) 
+$(op)(k, f::AbstractArray; kargs...) = $(op)(active_backend(), k, f::AbstractArray; kargs...) 
+function $(op)(b::AbstractNFFTBackend, k, f::AbstractArray; kargs...) 
   p = $(planfunc)(k, size(f); kargs... )
   return p * f
 end
+$(op)(::Missing, k, f::AbstractArray; kargs...) = no_backend_error()
 
-"""
-nfft_adjoint(k, N, fHat, rest...; kwargs...)
 
-calculates the adjoint nfft of the vector `fHat` for the nodes contained in the matrix `k`.
-The output is an array of size `N`
-"""
-function $(tfunc)(k, N, fHat;  kargs...) 
+$(tfunc)(k, N, fHat;  kargs...) = $(tfunc)(active_backend(), k, N, fHat;  kargs...)
+function $(tfunc)(b::AbstractNFFTBackend, k, N, fHat;  kargs...) 
   p = $(planfunc)(k, N;  kargs...)
   return $(trans)(p) * fHat
 end
+$(tfunc)(::Missing, k, N, fHat;  kargs...) = no_backend_error()
+
 
 end
 end

AbstractNFFTs/src/interface.jl

@@ -1,3 +1,33 @@
+abstract type AbstractNFFTBackend end
+struct BackendReference
+    ref::Ref{Union{Missing, AbstractNFFTBackend}}
+    BackendReference(val::Union{Missing, AbstractNFFTBackend}) = new(Ref{Union{Missing, AbstractNFFTBackend}}(val))
+end
+Base.setindex!(ref::BackendReference, val::Union{Missing, AbstractNFFTBackend}) = ref.ref[] = val
+Base.setindex!(ref::BackendReference, val::Module) = setindex!(ref, val.backend())
+Base.getindex(ref::BackendReference) = getindex(ref.ref)::Union{Missing, AbstractNFFTBackend}
+Base.convert(::Type{BackendReference}, val::AbstractNFFTBackend) = BackendReference(val)
+const nfft_backend = ScopedValue(BackendReference(missing))
+
+"""
+    set_active_backend!(back::Union{Missing, Module, AbstractNFFTBackend})
+
+Set the default NFFT plan backend. A module `back` must implement `back.backend()`.
+"""
+set_active_backend!(back::Module) = set_active_backend!(back.backend())
+function set_active_backend!(back::Union{Missing, AbstractNFFTBackend})
+  nfft_backend[][] = back
+end
+active_backend() = nfft_backend[][]
+function no_backend_error() 
+  error(
+    """
+    No default backend available!
+    Make sure to also "import/using" an NFFT backend such as NFFT or NonuniformFFTs.
+    """
+  )
+end
+
 """
   AbstractFTPlan{T,D,R}
 

NFFTTools/Project.toml

@@ -12,5 +12,5 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 [compat]
 julia = "1.6"
 AbstractFFTs = "1.0"
-AbstractNFFTs = "0.6, 0.7, 0.8"
+AbstractNFFTs = "0.6, 0.7, 0.8, 0.9"
 FFTW = "1"

Project.toml

@@ -1,7 +1,7 @@
 name = "NFFT"
 uuid = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 authors = ["Tobias Knopp <[email protected]>"]
-version = "0.13.7"
+version = "0.14"
 
 [deps]
 AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
@@ -19,7 +19,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
 Adapt = "3, 4"
-AbstractNFFTs = "0.8"
+AbstractNFFTs = "0.9"
 BasicInterpolators = "0.6.5, 0.7"
 DataFrames = "1.3.1, 1.4.1"
 FFTW = "1.5"