Define AbstractNFFTs.plan_nfft and create separate plan type (#42)

* Add wrapper type for AbstractNFFTs plans

* Pass kwargs to PlanNUFFT

* Define AbstractNFFTs.plan_nfft

* Update docs

* Update CHANGELOG + tests

* Try to fix tests on Julia 1.9

CHANGELOG.md

@@ -18,6 +18,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Currently, `FastApproximation()` is used on CPUs and `Direct()` on GPUs,
   where it is sometimes faster.
 
+- The `AbstractNFFTs.plan_nfft` function is now implemented for full compatibility with the AbstractNFFTs.jl interface.
+
 ### Changed
 
 - **BREAKING**: Change default precision of transforms.
@@ -26,6 +28,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Previously, the default was `m = HalfSupport(8)` and `σ = 2.0`, corresponding
   to a relative precision of the order of $10^{-14}$.
 
+- **BREAKING**: The `PlanNUFFT` constructor can no longer be used to create
+  plans compatible with AbstractNFFTs.jl / NFFT.jl.
+  Instead, a separate (and unexported) `NonuniformFFTs.NFFTPlan` type is now
+  defined which may be used for this purpose.
+  Alternatively, one can now use the `AbstractNFFTs.plan_nfft` function.
+
 - On GPUs, we now default to direct evaluation of kernel functions (e.g.
   Kaiser-Bessel) instead of polynomial approximations, as this seems to be
   faster and uses far fewer GPU registers.

docs/src/API.md

@@ -8,6 +8,7 @@ CurrentModule = NonuniformFFTs
 
 ```@docs
 PlanNUFFT
+NFFTPlan
 ```
 
 ## Setting non-uniform points

src/abstractNFFTs.jl

@@ -6,25 +6,92 @@
 # - it doesn't allow simultaneous (directional) transforms (can be fixed within NonuniformFFTs.jl)
 
 using LinearAlgebra: LinearAlgebra, Adjoint
+using Adapt: adapt
 
-AbstractNFFTs.size_in(p::PlanNUFFT) = size(p)  # array dimensions in uniform space
-AbstractNFFTs.size_out(p::PlanNUFFT) = (length(p.points),)
+# This is a wrapper type allowing to define an interface which is compatible with
+# AbstractNFFTs.jl. It is not exported to avoid clashes with NFFT.jl.
+"""
+    NonuniformFFTs.NFFTPlan(xp::AbstractMatrix{T}, dims::Dims{D}; kwargs...) -> plan
+
+Create a NUFFT plan which is compatible with the
+[AbstractNFFTs.jl](https://juliamath.github.io/NFFT.jl/stable/abstract/) interface.
+
+The plan follows different rules from plans created by [`PlanNUFFT`](@ref).
+In particular:
+
+- points are assumed to be in ``[-1/2, 1/2)`` instead of ``[0, 2π)``;
+- the opposite Fourier sign convention is used (e.g. ``e^{-i k x_j}`` becomes ``e^{+2π i k x_j}``);
+- uniform data is in increasing order by default, with frequencies ``k = -N/2, …, -1, 0,
+  1, …, N/2-1``, as opposed to preserving the order used by FFTW (which starts at ``k = 0``);
+- it only supports complex non-uniform data.
+
+This constructor requires passing the non-uniform locations `xp` as the first argument.
+These should be given as a matrix of dimensions `(D, Np)`, where `D` is the spatial
+dimension and `Np` the number of non-uniform points.
+
+The second argument is simply the size `(N₁, N₂, …)` of the uniform data arrays.
+
+Most keyword arguments from [`PlanNUFFT`](@ref) are also accepted here.
+Moreover, for compatibility reasons, most keyword arguments from the NFFT.jl package are
+also accepted as detailed below.
+
+This type of plan can also be created via the
+[`AbstractNFFTs.plan_nfft`](https://juliamath.github.io/NFFT.jl/v0.13.5/abstract/#Plan-Interface)
+function.
+
+This constructor creates a plan which assumes complex-valued non-uniform data.
+For real-valued data, the [`PlanNUFFT`](@ref) constructor should be used instead.
+
+# Compatibility with NFFT.jl
+
+Most of the [parameters](https://juliamath.github.io/NFFT.jl/stable/overview/#Parameters)
+supported by the NFFT.jl package are also supported by this constructor.
+The currently supported parameters are `reltol`, `m`, `σ`, `window`, `blocking`, `sortNodes` and `fftflags`.
+
+Moreover, unlike [`PlanNUFFT`](@ref), this constructor sets `fftshift = true` by default (but
+can be overridden) so that the uniform data ordering is the same as in NFFT.jl.
+
+!!! warning "Type instability"
+
+    Explicitly passing some of these parameters may result in type-unstable code, since the
+    exact type of the returned plan cannot be inferred.
+    This is because, in NonuniformFFTs.jl, parameters such as the kernel size (`m`) or the
+    convolution window (`window`) are included in the plan type (they are compile-time constants).
+    
+# GPU usage
+
+To create a GPU-compatible plan, simply pass the locations `xp` as a GPU array (e.g. a `CuArray` in CUDA).
+Unlike [`PlanNUFFT`](@ref), the `backend` argument is not needed here and will be simply ignored.
+"""
+struct NFFTPlan{
+        T <: AbstractFloat, N, Plan <: PlanNUFFT{Complex{T}, N, 1},
+    } <: AbstractNFFTPlan{T, N, 1}
+    p :: Plan
+end
+
+function Base.show(io::IO, p::NFFTPlan{T, N}) where {T, N}
+    print(io, "NonuniformFFTs.NFFTPlan{$T, $N} wrapping a PlanNUFFT:\n")
+    print(io, p.p)
+end
+
+AbstractNFFTs.size_in(p::NFFTPlan) = size(p.p)  # array dimensions in uniform space
+AbstractNFFTs.size_out(p::NFFTPlan) = (length(p.p.points),)
 
 # Uniform to non-uniform
 function LinearAlgebra.mul!(
-        vp::AbstractVector{Tc}, p::PlanNUFFT{Tc}, ûs::AbstractArray{Tc},
-    ) where {Tc <: Complex}
-    exec_type2!(vp, p, ûs)
+        vp::AbstractVector{Complex{T}}, p::NFFTPlan{T}, ûs::AbstractArray{Complex{T}},
+    ) where {T}
+    exec_type2!(vp, p.p, ûs)
     vp
 end
 
 # Non-uniform to uniform
 function LinearAlgebra.mul!(
-        ûs::AbstractArray{Tc, N},
-        p::Adjoint{Tc, <:PlanNUFFT{Tc, N}},
-        vp::AbstractVector{Tc},
-    ) where {Tc <: Complex, N}
-    exec_type1!(ûs, parent(p), vp)
+        ûs::AbstractArray{Complex{T}, N},
+        p::Adjoint{Complex{T}, <:NFFTPlan{T, N}},
+        vp::AbstractVector{Complex{T}},
+    ) where {T, N}
+    exec_type1!(ûs, parent(p).p, vp)
     vp
 end
 
@@ -44,14 +111,14 @@ end
 # Takes locations in [-1/2, 1/2)ᵈ, so we need to transform the point convention.
 # Note: the `transform` argument is an internal of set_points! and is not documented.
 # It takes either an NTuple (a point location x⃗) or a scalar (coordinate xᵢ) as input.
-function AbstractNFFTs.nodes!(p::PlanNUFFT, xp::AbstractMatrix{T}) where {T <: AbstractFloat}
-    set_points!(p, xp; transform = _transform_point_convention)
+function AbstractNFFTs.nodes!(p::NFFTPlan, xp::AbstractMatrix{T}) where {T <: AbstractFloat}
+    set_points!(p.p, xp; transform = _transform_point_convention)
 end
 
 # This is to avoid ambiguity issues, since AbstractNFFTs defines nodes! for Matrix instead
 # of AbstractMatrix.
-function AbstractNFFTs.nodes!(p::PlanNUFFT{Complex{T}}, xp::Matrix{T}) where {T <: AbstractFloat}
-    invoke(AbstractNFFTs.nodes!, Tuple{PlanNUFFT, AbstractMatrix{T}}, p, xp)
+function AbstractNFFTs.nodes!(p::NFFTPlan{T}, xp::Matrix{T}) where {T <: AbstractFloat}
+    invoke(AbstractNFFTs.nodes!, Tuple{NFFTPlan, AbstractMatrix{T}}, p, xp)
 end
 
 Base.@constprop :aggressive function convert_window_function(w::Symbol, backend)
@@ -79,27 +146,45 @@ end
 # We use @constprop to hopefully avoid some type instabilities, but that doesn't seem to
 # fully help here.
 # TODO: support all keyword arguments of the standard constructor
-Base.@constprop :aggressive function PlanNUFFT(
-        xp::AbstractMatrix{Tr}, Ns::Dims;
+Base.@constprop :aggressive function NFFTPlan(
+        xp::AbstractMatrix{T}, Ns::Dims;
         fftflags = FFTW.ESTIMATE, blocking = true, sortNodes = false,
         window = default_kernel(KA.get_backend(xp)),
         fftshift = true,  # for compatibility with NFFT.jl
-        synchronise = false,
-        kwargs...,
-    ) where {Tr <: AbstractFloat}
+        kws...,
+    ) where {T <: AbstractFloat}
     # Note: the NFFT.jl package uses an odd window size, w = 2m + 1.
     # Here we use an even window size w = 2m, which should result in a slightly lower
     # accuracy for the same m.
-    m, σ, reltol = AbstractNFFTs.accuracyParams(; kwargs...)
+    kws_plan, kws_accuracy = _split_accuracy_params(; kws...)
+    m_actual, σ_actual, reltol_actual = AbstractNFFTs.accuracyParams(; kws_accuracy...)
     backend = KA.get_backend(xp)  # e.g. use GPU backend if xp is a GPU array
     sort_points = sortNodes ? True() : False()  # this is type-unstable (unless constant propagation happens)
     block_size = blocking ? default_block_size(Ns, backend) : nothing  # also type-unstable
     kernel = window isa AbstractKernel ? window : convert_window_function(window, backend)
+    get(kws, :backend, backend) == backend || throw(ArgumentError("`backend` argument is incompatible with array type"))
     p = PlanNUFFT(
-        Complex{Tr}, Ns, HalfSupport(m);
-        backend, σ = Tr(σ), sort_points, fftshift, block_size,
-        kernel, fftw_flags = fftflags, synchronise,
+        Complex{T}, Ns, HalfSupport(m_actual);
+        backend, σ = T(σ_actual), sort_points, fftshift, block_size,
+        kernel, fftw_flags = fftflags,
+        kws_plan...,
     )
-    AbstractNFFTs.nodes!(p, xp)
-    p
+    pp = NFFTPlan(p)
+    AbstractNFFTs.nodes!(pp, xp)
+    pp
+end
+
+function _split_accuracy_params(; kws...)
+    nt = NamedTuple(kws)
+    nt_plan = Base.structdiff(nt, NamedTuple{(:m, :σ, :reltol)})  # remove accuracy params
+    nt_accuracy = Base.structdiff(nt, nt_plan)  # only accuracy params
+    nt_plan, nt_accuracy
+end
+
+function AbstractNFFTs.plan_nfft(
+        ::Type{Q}, xp::AbstractMatrix{T}, Ns::Dims{D};
+        kwargs...,
+    ) where {Q, T, D}
+    xp_q = adapt(Q, xp)  # convert array (and copy to device) if needed
+    NFFTPlan(xp_q, Ns; kwargs...)
 end

src/plan.jl

@@ -132,7 +132,7 @@ the order of ``10^{-7}`` for `Float64` or `ComplexF64` data.
 
 ## Other performance parameters
 
-These are more advanced performance parameters which may dissappear or whose behaviour may
+These are more advanced performance parameters which may disappear or whose behaviour may
 change in the future.
 
 - `sort_points = False()`: whether to internally permute the order of the non-uniform points.
@@ -201,43 +201,6 @@ since the size of the first dimension is divided roughly by 2 (taking advantage
 symmetry).
 For convenience, one can call [`size(::PlanNUFFT)`](@ref) on the constructed plan to know in
 advance the dimensions of the uniform data arrays.
-
----
-
-    PlanNUFFT(xp::AbstractMatrix{T}, dims::Dims{D}; kwargs...)
-
-Create a [`PlanNUFFT`](@ref) which is compatible with the
-[AbstractNFFTs.jl](https://juliamath.github.io/NFFT.jl/stable/abstract/) interface.
-
-This constructor requires passing the non-uniform locations `xp` as the first argument.
-These should be given as a matrix of dimensions `(D, Np)`, where `D` is the spatial
-dimension and `Np` the number of non-uniform points.
-
-The second argument is simply the size `(N₁, N₂, …)` of the uniform data arrays.
-
-This variant creates a plan which assumes complex-valued non-uniform data.
-For real-valued data, the other constructor should be used instead.
-
-# Compatibility with NFFT.jl
-
-Most of the [parameters](https://juliamath.github.io/NFFT.jl/stable/overview/#Parameters)
-supported by the NFFT.jl package are also supported by this constructor.
-The currently supported parameters are `reltol`, `m`, `σ`, `window`, `blocking`, `sortNodes` and `fftflags`.
-
-Moreover, unlike the first variant, this constructor sets `fftshift = true` by default (but
-can be overridden) so that the uniform data ordering is the same as in NFFT.jl.
-
-!!! warning "Type instability"
-
-    Explicitly passing some of these parameters may result in type-unstable code, since the
-    exact type of the returned plan cannot be inferred.
-    This is because, in NonuniformFFTs.jl, parameters such as the kernel size (`m`) or the
-    convolution window (`window`) are included in the plan type (they are compile-time constants).
-
-# GPU usage
-
-To create a GPU-compatible plan, simply pass the locations `xp` as a GPU array (e.g. a `CuArray` in CUDA).
-Unlike the first constructor, the `backend` argument is not needed here and will be simply ignored.
 """
 struct PlanNUFFT{
         T <: Number,  # non-uniform data type (can be real or complex)
@@ -253,7 +216,7 @@ struct PlanNUFFT{
         Blocks <: AbstractBlockData,
         IndexMap <: NTuple{N, AbstractVector{Int}},
         Timer <: TimerOutput,
-    } <: AbstractNFFTPlan{Treal, N, 1}  # the AbstractNFFTPlan only really makes sense when T <: Complex
+    }
     kernels :: Kernels
     backend :: Backend  # CPU, GPU, ...
     kernel_evalmode :: KernelEvalMode
@@ -267,6 +230,10 @@ struct PlanNUFFT{
     synchronise :: Bool
 end
 
+# This represents the type of data in Fourier space.
+# It is defined like this for compatibility with AbstractNFFTs plans.
+Base.eltype(::PlanNUFFT{T}) where {T} = complex(T)
+
 function Base.show(io::IO, p::PlanNUFFT{T, N, Nc}) where {T, N, Nc}
     (; kernels, backend, σ, blocks, fftshift,) = p
     M = Kernels.half_support(first(kernels))

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"

test/abstractNFFTs.jl

@@ -15,10 +15,10 @@ function compare_with_nfft(Ns; Np = 1000)
     reltol = 1e-9
     window = :kaiser_bessel
 
-    p = PlanNUFFT(xp, Ns; reltol, window,)
+    p = NonuniformFFTs.NFFTPlan(xp, Ns; reltol, window,)
     p_nfft = NFFT.NFFTPlan(xp, Ns; reltol, window,)
 
-    @test startswith(repr(p), "$(d)-dimensional PlanNUFFT with input type $(Complex{T}):")  # test pretty-printing
+    @test startswith(repr(p), "NonuniformFFTs.NFFTPlan{$T, $d} wrapping a PlanNUFFT:")  # test pretty-printing
 
     @test size_in(p) === size_in(p_nfft)
     @test size_out(p) === size_out(p_nfft)

test/pseudo_gpu.jl

@@ -12,6 +12,7 @@
 using NonuniformFFTs
 using StaticArrays: SVector  # for convenience
 using KernelAbstractions: KernelAbstractions as KA
+using AbstractNFFTs: AbstractNFFTs
 using Adapt: Adapt, adapt
 using GPUArraysCore: AbstractGPUArray
 using Random
@@ -75,7 +76,7 @@ function run_plan(p::PlanNUFFT, xp_init::AbstractArray, vp_init::NTuple{Nc, Abst
         end
     end
 
-    T = eltype(p)  # this is actually defined in AbstractNFFTs; it represents the type in Fourier space (always complex)
+    T = eltype(p)  # type in Fourier space (always complex) - for compatibility with AbstractNFFTs plans
     @test T <: Complex
     dims = size(p)
     us = map(_ -> KA.allocate(backend, T, dims), vp)
@@ -107,6 +108,16 @@ function compare_with_cpu(::Type{T}, dims; Np = prod(dims), ntransforms::Val{Nc}
     p_cpu = @inferred PlanNUFFT(T, dims; params..., backend = CPU())
     p_gpu = @inferred PlanNUFFT(T, dims; params..., backend = PseudoGPU())
 
+    # Test that plan_nfft interface works.
+    @testset "AbstractNFFTs.plan_nfft" begin
+        xmat = reinterpret(reshape, Tr, xp_init)
+        p_nfft = @inferred AbstractNFFTs.plan_nfft(PseudoGPUArray, xmat, dims)
+        @test p_nfft.p.backend isa PseudoGPU
+        # Test without the initial argument (type of array)
+        p_nfft_cpu = @inferred AbstractNFFTs.plan_nfft(xmat, dims)
+        @test p_nfft_cpu.p.backend isa CPU
+    end
+
     r_cpu = run_plan(p_cpu, xp_init, vp_init)
     r_gpu = run_plan(p_gpu, xp_init, vp_init)