update nufft

reverse type I and II to match the paper

add `lambertw` as a temporary measure

add docstrings and section in readme

Author: MikaelSlevinsky

README.md

@@ -75,6 +75,34 @@ is valid for the half-open square `(α,β) ∈ (-1/2,1/2]^2`. Therefore, the fas
 when the parameters are inside. If the parameters `(α,β)` are not exceptionally beyond the square,
 then increment/decrement operators are used with linear complexity (and linear conditioning) in the degree.
 
+## Nonuniform fast Fourier transforms
+
+The NUFFTs are implemented thanks to [Alex Townsend](https://github.com/ajt60gaibb). `nufft1` assumes uniform samples and noninteger frequencies, while `nufft2` assumes nonuniform samples and integer frequencies.
+```julia
+julia> n = 10^4;
+
+julia> c = complex(rand(n));
+
+julia> ω = collect(0:n-1) + rand(n);
+
+julia> nufft1(c, ω, eps());
+
+julia> p1 = plan_nufft1(ω, eps());
+
+julia> @time p1*c;
+  0.002383 seconds (6 allocations: 156.484 KiB)
+
+julia> x = (collect(0:n-1) + 3rand(n))/n;
+
+julia> nufft2(c, x, eps());
+
+julia> p2 = plan_nufft2(x, eps());
+
+julia> @time p2*c;
+  0.001478 seconds (6 allocations: 156.484 KiB)
+
+```
+
 ## The Padua Transform
 
 The Padua transform and its inverse are implemented thanks to [Michael Clarke](https://github.com/MikeAClarke). These are optimized methods designed for computing the bivariate Chebyshev coefficients by interpolating a bivariate function at the Padua points on `[-1,1]^2`.
@@ -87,7 +115,7 @@ julia> N = div((n+1)*(n+2),2);
 julia> v = rand(N); # The length of v is the number of Padua points
 
 julia> @time norm(ipaduatransform(paduatransform(v))-v)
-0.006571 seconds (846 allocations: 1.746 MiB)
+  0.006571 seconds (846 allocations: 1.746 MiB)
 3.123637691861415e-14
 
 ```
@@ -126,10 +154,12 @@ As with other fast transforms, `plan_sph2fourier` saves effort by caching the pr
 
    [2]  N. Hale and A. Townsend. <a href="http://dx.doi.org/10.1137/130932223">A fast, simple, and stable Chebyshev—Legendre transform using and asymptotic formula</a>, *SIAM J. Sci. Comput.*, **36**:A148—A167, 2014.
 
-   [3] J. Keiner. <a href="http://dx.doi.org/10.1137/070703065">Computing with expansions in Gegenbauer polynomials</a>, *SIAM J. Sci. Comput.*, **31**:2151—2171, 2009.
+   [3]  J. Keiner. <a href="http://dx.doi.org/10.1137/070703065">Computing with expansions in Gegenbauer polynomials</a>, *SIAM J. Sci. Comput.*, **31**:2151—2171, 2009.
+
+   [4]  D. Ruiz—Antolín and A. Townsend. <a href="https://arxiv.org/abs/1701.04492">A nonuniform fast Fourier transform based on low rank approximation</a>, arXiv:1701.04492, 2017.
 
-   [4]  R. M. Slevinsky. <a href="https://doi.org/10.1093/imanum/drw070">On the use of Hahn's asymptotic formula and stabilized recurrence for a fast, simple, and stable Chebyshev—Jacobi transform</a>, in press at *IMA J. Numer. Anal.*, 2017.
+   [5]  R. M. Slevinsky. <a href="https://doi.org/10.1093/imanum/drw070">On the use of Hahn's asymptotic formula and stabilized recurrence for a fast, simple, and stable Chebyshev—Jacobi transform</a>, in press at *IMA J. Numer. Anal.*, 2017.
 
-   [5]  R. M. Slevinsky. <a href="https://arxiv.org/abs/1705.05448">Fast and backward stable transforms between spherical harmonic expansions and bivariate Fourier series</a>, arXiv:1705.05448, 2017.
+   [6]  R. M. Slevinsky. <a href="https://arxiv.org/abs/1705.05448">Fast and backward stable transforms between spherical harmonic expansions and bivariate Fourier series</a>, arXiv:1705.05448, 2017.
 
-   [6]  A. Townsend, M. Webb, and S. Olver. <a href="https://doi.org/10.1090/mcom/3277">Fast polynomial transforms based on Toeplitz and Hankel matrices</a>, in press at *Math. Comp.*, 2017.
+   [7]  A. Townsend, M. Webb, and S. Olver. <a href="https://doi.org/10.1090/mcom/3277">Fast polynomial transforms based on Toeplitz and Hankel matrices</a>, in press at *Math. Comp.*, 2017.

docs/src/index.md

@@ -52,6 +52,22 @@ plan_cjt
 plan_icjt
 ```
 
+```@docs
+nufft1
+```
+
+```@docs
+nufft2
+```
+
+```@docs
+plan_nufft1
+```
+
+```@docs
+plan_nufft2
+```
+
 ```@docs
 paduatransform
 ```
@@ -112,6 +128,10 @@ FastTransforms.δ
 FastTransforms.Λ
 ```
 
+```@docs
+FastTransforms.lambertw
+```
+
 ```@docs
 FastTransforms.pochhammer
 ```

src/FastTransforms.jl

@@ -16,11 +16,14 @@ export normleg2cheb, cheb2normleg, normleg12cheb2, cheb22normleg1
 export plan_leg2cheb, plan_cheb2leg
 export plan_normleg2cheb, plan_cheb2normleg
 export plan_normleg12cheb2, plan_cheb22normleg1
+
 export gaunt
+
 export paduatransform, ipaduatransform, paduatransform!, ipaduatransform!, paduapoints
 export plan_paduatransform!, plan_ipaduatransform!
+
 export nufft, nufft1, nufft2
-export nufft_plan, nufft1_plan, nufft2_plan
+export plan_nufft, plan_nufft1, plan_nufft2
 
 export SlowSphericalHarmonicPlan, FastSphericalHarmonicPlan, ThinSphericalHarmonicPlan
 export sph2fourier, fourier2sph, plan_sph2fourier

src/nufft.jl

@@ -1,93 +1,185 @@
-function nufft1_plan{T<:AbstractFloat}( x::AbstractVector{T}, ϵ::T )
-
-t_idx = AssignClosestEquispacedFFTpoint( x )
-γ = PerturbationParameter( x, AssignClosestEquispacedGridpoint( x ) )
-K = FindK(γ, ϵ)   
-u = constructU(x, K)
-v = constructV(x, K)
-p( c ) = (u.*(fft(Diagonal(c)*v,1)[t_idx,:]))*ones(K)
+doc"""
+Pre-compute a nonuniform fast Fourier transform of type `N`.
+
+For best performance, choose the right number of threads by `FFTW.set_num_threads(4)`, for example.
+"""
+immutable NUFFTPlan{N,T,FFT} <: Base.DFT.Plan{T}
+    U::Matrix{T}
+    V::Matrix{T}
+    p::FFT
+    t::Vector{Int}
+    temp::Matrix{T}
+    temp2::Matrix{T}
+    Ones::Vector{T}
 end
 
-function nufft2_plan{T<:AbstractFloat}( ω::AbstractVector{T}, ϵ::T )
+doc"""
+Computes a nonuniform fast Fourier transform of type I:
+
+```math
+f_j = \sum_{k=1}^N c_k e^{-2\pi{\rm i} (j-1)/N \omega_k},\quad{\rm for}\quad 1 \le j \le N.
+```
+"""
+function plan_nufft1{T<:AbstractFloat}(ω::AbstractVector{T}, ϵ::T)
+    N = length(ω)
+    ωdN = ω/N
+    t = AssignClosestEquispacedFFTpoint(ωdN)
+    γ = PerturbationParameter(ωdN, AssignClosestEquispacedGridpoint(ωdN))
+    K = FindK(γ, ϵ)
+    U = constructU( ωdN, K)
+    V = constructV( ωdN, K)
+    p = plan_ifft!(V, 1)
+    temp = zeros(Complex{T}, N, K)
+    temp2 = zeros(Complex{T}, N, K)
+    Ones = ones(Complex{T}, K)
 
-N = size(ω, 1)
-t_idx = AssignClosestEquispacedFFTpoint( ω/N )
-γ = PerturbationParameter( ω/N, AssignClosestEquispacedGridpoint( ω/N ) )
-K = FindK(γ, ϵ) 
-u = constructU( ω/N, K)
-v = constructV( ω/N, K) 
-In = speye(Complex{T},  N, N)
-p( c ) = (v.*(N*conj(ifft(In[:,t_idx]*conj(Diagonal(c)*u),1))))*ones(K)
+    NUFFTPlan{1, eltype(U), typeof(p)}(U, V, p, t, temp, temp2, Ones)
 end
 
-nufft_plan{T<:AbstractFloat}( x::AbstractVector{T}, ϵ::T ) = nufft1_plan( x, ϵ )
-nufft{T<:AbstractFloat}( c::AbstractVector, x::AbstractVector{T}, ϵ::T ) = nufft_plan(x, ϵ)(c) 
-nufft1{T<:AbstractFloat}( c::AbstractVector, x::AbstractVector{T}, ϵ::T ) = nufft1_plan(x, ϵ)(c) 
-nufft2{T<:AbstractFloat}( c::AbstractVector, ω::AbstractVector{T}, ϵ::T ) = nufft2_plan(ω, ϵ)(c)
+doc"""
+Computes a nonuniform fast Fourier transform of type II:
 
-FindK{T<:AbstractFloat}(γ::T, ϵ::T) = Int( ceil(5.0*γ.*exp(lambertw(log(10.0/ϵ)./γ/7.0))) )
-AssignClosestEquispacedGridpoint{T<:AbstractFloat}( x::AbstractVector{T} )::AbstractVector{T} = round(size(x,1)*x)
-AssignClosestEquispacedFFTpoint{T<:AbstractFloat}( x::AbstractVector{T} )::Array{Int64,1} = mod(round(Int64, size(x,1)*x), size(x,1)) + 1
-PerturbationParameter{T<:AbstractFloat}( x::AbstractVector{T}, s_vec::AbstractVector{T} )::AbstractFloat = norm( size(x,1)*x - s_vec, Inf)
+```math
+f_j = \sum_{k=1}^N c_k e^{-2\pi{\rm i} x_j (k-1)},\quad{\rm for}\quad 1 \le j \le N.
+```
+"""
+function plan_nufft2{T<:AbstractFloat}(x::AbstractVector{T}, ϵ::T)
+    N = length(x)
+    t = AssignClosestEquispacedFFTpoint(x)
+    γ = PerturbationParameter(x, AssignClosestEquispacedGridpoint(x))
+    K = FindK(γ, ϵ)
+    U = constructU(x, K)
+    V = constructV(x, K)
+    p = plan_fft!(U, 1)
+    temp = zeros(Complex{T}, N, K)
+    temp2 = zeros(Complex{T}, N, K)
+    Ones = ones(Complex{T}, K)
+
+    NUFFTPlan{2, eltype(U), typeof(p)}(U, V, p, t, temp, temp2, Ones)
+end
 
-function constructU{T<:AbstractFloat}(x::AbstractVector{T}, K::Int64) 
-# Construct a low rank approximation, using Chebyshev expansions 
-# for AK = exp(-2*pi*1im*(x[j]-j/N)*k): 
+function (*){N,T,V}(p::NUFFTPlan{N,T}, x::AbstractVector{V})
+    A_mul_B!(zeros(promote_type(T,V), length(x)), p, x)
+end
 
-N = size(x, 1)
-#(s_vec, t_idx, γ) = FindAlgorithmicParameters( x ) 
-s_vec = AssignClosestEquispacedGridpoint( x )
-er = N*x - s_vec
-γ = norm( er, Inf )
+function Base.A_mul_B!{T}(y::AbstractVector{T}, P::NUFFTPlan{1,T}, c::AbstractVector{T})
+    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones
 
-# colspace vectors:
-u = Diagonal(exp(-1im*pi*er))*ChebyshevP(K-1, er/γ)*Bessel_coeffs(K, γ)
-end 
+    # (V.*(N*conj(ifft(In[:,t]*conj(Diagonal(c)*U),1))))*ones(K)
 
-function constructV{T<:AbstractFloat}(x::AbstractVector{T}, K::Int64)
+    broadcast!(*, temp, c, U)
+    conj!(temp)
+    fill!(temp2, zero(T))
+    recombine_rows!(temp, t, temp2)
+    p*temp2
+    conj!(temp2)
+    broadcast!(*, temp, V, temp2)
+    A_mul_B!(y, temp, Ones)
+    scale!(length(c), y)
 
-N = size(x, 1)
-v = complex(ChebyshevP(K-1, 2.0*collect(0:N-1)/N - ones(N) ))
+    y
 end
 
-function Bessel_coeffs{T<:AbstractFloat}(K::Int64, γ::T)::Array{Complex{T},2}
-# Calculate the Chebyshev coefficients of exp(-2*pi*1im*x*y) on [-gam,gam]x[0,1]
+function Base.A_mul_B!{T}(y::AbstractVector{T}, P::NUFFTPlan{2,T}, c::AbstractVector{T})
+    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones
+
+    # (U.*(fft(Diagonal(c)*V,1)[t+1,:]))*ones(K)
+
+    broadcast!(*, temp, c, V)
+    p*temp
+    reindex_temp!(temp, t, temp2)
+    broadcast!(*, temp, U, temp2)
+    A_mul_B!(y, temp, Ones)
+
+    y
+end
 
-cfs = complex(zeros( K, K ))
-arg = -γ*pi/2.0
-for p = 0:K-1
- 	for q = mod(p,2):2:K-1 
-		cfs[p+1,q+1] = 4.0*(1im)^q*besselj((p+q)/2,arg).*besselj((q-p)/2,arg)
-	end 
+function reindex_temp!{T}(temp::Matrix{T}, t::Vector{Int}, temp2::Matrix{T})
+    @inbounds for j = 1:size(temp, 2)
+        for i = 1:size(temp, 1)
+            temp2[i, j] = temp[t[i], j]
+        end
+    end
+    temp2
 end
-cfs[1,:] = cfs[1,:]/2.0
-cfs[:,1] = cfs[:,1]/2.0
-return cfs
+
+function recombine_rows!{T}(temp::Matrix{T}, t::Vector{Int}, temp2::Matrix{T})
+    @inbounds for j = 1:size(temp, 2)
+        for i = 1:size(temp, 1)
+            temp2[t[i], j] += temp[i, j]
+        end
+    end
+    temp2
 end
 
-function ChebyshevP{T<:AbstractFloat}(n::Int64, x::AbstractVector{T})::AbstractArray{T} 
-# Evaluate Chebyshev polynomials of degree 0,...,n at x:
+doc"""
+Pre-compute a nonuniform fast Fourier transform of type I.
+"""
+nufft1{T<:AbstractFloat}(c::AbstractVector, ω::AbstractVector{T}, ϵ::T) = plan_nufft1(ω, ϵ)*c
+
+doc"""
+Pre-compute a nonuniform fast Fourier transform of type II.
+"""
+nufft2{T<:AbstractFloat}(c::AbstractVector, x::AbstractVector{T}, ϵ::T) = plan_nufft2(x, ϵ)*c
+
+FindK{T<:AbstractFloat}(γ::T, ϵ::T) = Int(ceil(5*γ*exp(lambertw(log(10/ϵ)/γ/7))))
+AssignClosestEquispacedGridpoint{T<:AbstractFloat}(x::AbstractVector{T})::AbstractVector{T} = round.([Int], size(x, 1)*x)
+AssignClosestEquispacedFFTpoint{T<:AbstractFloat}(x::AbstractVector{T})::Array{Int,1} = mod.(round.([Int], size(x, 1)*x), size(x, 1)) + 1
+PerturbationParameter{T<:AbstractFloat}(x::AbstractVector{T}, s_vec::AbstractVector{T})::AbstractFloat = norm(size(x, 1)*x - s_vec, Inf)
 
-N = size(x, 1)
-Tcheb = Array{T}(N, n+1)
+function constructU{T<:AbstractFloat}(x::AbstractVector{T}, K::Int)
+    # Construct a low rank approximation, using Chebyshev expansions
+    # for AK = exp(-2*pi*1im*(x[j]-j/N)*k):
+    N = size(x, 1)
+    #(s_vec, t, γ) = FindAlgorithmicParameters( x )
+    s_vec = AssignClosestEquispacedGridpoint(x)
+    er = N*x - s_vec
+    γ = norm(er, Inf)
+    # colspace vectors:
+    Diagonal(exp.(-im*(pi*er)))*ChebyshevP(K-1, er/γ)*Bessel_coeffs(K, γ)
+end
 
-# T_0(x) = 1.0
-One = convert(eltype(x),1.0)
-@inbounds for j = 1:N
-	Tcheb[j, 1] = One
+function constructV{T<:AbstractFloat}(x::AbstractVector{T}, K::Int)
+    N = size(x, 1)
+    complex(ChebyshevP(K-1, two(T)*collect(0:N-1)/N - ones(N) ))
 end
-# T_1(x) = x
-if ( n > 0 ) 
-	@inbounds for j = 1:N
-		Tcheb[j, 2] = x[j]
-	end
+
+function Bessel_coeffs{T<:AbstractFloat}(K::Int, γ::T)
+    # Calculate the Chebyshev coefficients of exp(-2*pi*1im*x*y) on [-gam,gam]x[0,1]
+    cfs = zeros(Complex{T}, K, K)
+    arg = -γ*π/two(T)
+    for p = 0:K-1
+     	for q = mod(p,2):2:K-1
+    		cfs[p+1,q+1] = 4*(1im)^q*besselj((p+q)/2,arg).*besselj((q-p)/2,arg)
+    	end
+    end
+    cfs[1,:] = cfs[1,:]/two(T)
+    cfs[:,1] = cfs[:,1]/two(T)
+    return cfs
 end
-# 3-term recurrence relation: 
-twoX = 2x
-@inbounds for k = 2:n
-	@inbounds for j = 1:N
-    		Tcheb[j, k+1] = twoX[j]*Tcheb[j, k] - Tcheb[j, k-1]
-	end
+
+function ChebyshevP{T<:AbstractFloat}(n::Int, x::AbstractVector{T})
+    # Evaluate Chebyshev polynomials of degree 0,...,n at x:
+    N = size(x, 1)
+    Tcheb = Matrix{T}(N, n+1)
+
+    # T_0(x) = 1.0
+    One = convert(eltype(x),1.0)
+    @inbounds for j = 1:N
+        Tcheb[j, 1] = One
+    end
+    # T_1(x) = x
+    if ( n > 0 )
+        @inbounds for j = 1:N
+            Tcheb[j, 2] = x[j]
+        end
+    end
+    # 3-term recurrence relation:
+    twoX = 2x
+    @inbounds for k = 2:n
+        @simd for j = 1:N
+            Tcheb[j, k+1] = twoX[j]*Tcheb[j, k] - Tcheb[j, k-1]
+        end
+    end
+    return Tcheb
 end
-return Tcheb
-end