fix and improve fft methods (#804)

Author: jishnub

src/Extras/fftBigFloat.jl

@@ -1,17 +1,21 @@
 # This is a Cooley-Tukey FFT algorithm for any number type.
-function fft_pow2(x::Vector{F}) where F
+function fft_pow2(x::StridedVector)
     n = length(x)
-    T = mapreduce(eltype,promote_type,x)
+    T = mapreduce(eltype(x) isa Fun ? cfstype : eltype, promote_type,x)
     @assert ispow2(n)
     if n==1
-        return x
+        return convert(Vector, x)
     elseif n==2
-        return F[x[1]+x[2];x[1]-x[2]]
+        return eltype(x)[x[1]+x[2], x[1]-x[2]]
     end
-    even,odd = fft_pow2(x[1:2:end-1]),fft_pow2(x[2:2:end])
-    twiddle = exp(-2im*convert(T,π)/n*collect(0:n-1))
-    half1 = even + odd.*twiddle[1:div(n,2)]
-    half2 = even + odd.*twiddle[div(n,2)+1:n]
-    return vcat(half1,half2)
+    even,odd = fft_pow2(@view x[1:2:end-1]), fft_pow2(@view x[2:2:end])
+    ret = similar(x, n)
+    for (halfind, ind) in enumerate(1:div(n,2))
+        ret[ind] = even[halfind] + odd[halfind] * cis(-2convert(T,π)/n * (ind - 1))
+    end
+    for (halfind, ind) in enumerate(div(n,2)+1:n)
+        ret[ind] = even[halfind] + odd[halfind] * cis(-2convert(T,π)/n * (ind - 1))
+    end
+    return ret
 end
-ifft_pow2(x::Vector) = conj(fft_pow2(conj(x)))/length(x)
+ifft_pow2(x::StridedVector) = conj(fft_pow2(conj(x)))/length(x)

src/Extras/fftGeneric.jl

@@ -1,30 +1,29 @@
-const BigFloats = Union{BigFloat,Complex{BigFloat}}
-
-function fft(x::AbstractVector{F}) where F<:Fun
-    n,T = length(x),mapreduce(eltype,promote_type,x)
+function fft(x::AbstractVector{<:Fun})
+    n,T = length(x),mapreduce(cfstype,promote_type,x)
     if ispow2(n) return fft_pow2(x) end
     ks = range(zero(real(T)),stop=n-one(real(T)),length=n)
-    Wks = exp(-im*convert(T,π)*ks.^2/n)
-    xq,wq = x.*Wks,conj([exp(-im*convert(T,π)*n);reverse(Wks);Wks[2:end]])
-    return Wks.*conv(xq,wq)[n+1:2n]
+    Wks = cis.(-convert(T,π)/n .* ks.^2)
+    xq,wq = x.*Wks,conj([cis(-convert(T,π)*n);reverse(Wks); @view Wks[2:end]])
+    return Wks.* @view conv(xq,wq)[n+1:2n]
 end
 
-ifft(x::AbstractVector{F}) where {F<:Fun} = conj(fft(conj(x)))/length(x)
-function ifft!(x::AbstractVector{F}) where F<:Fun
+ifft(x::AbstractVector{<:Fun}) = conj(fft(conj(x)))/length(x)
+function ifft!(x::AbstractVector{<:Fun})
     y = conj(fft(conj(x)))/length(x)
-    x[:] = y
+    x .= y
     return x
 end
 
-function conv(u::StridedVector{F}, v::StridedVector) where F<:Fun
+nextpow2(n) = 2^ceil(Int, Base.log2(n))
+function conv(u::StridedVector{<:Fun}, v::StridedVector)
     nu,nv = length(u),length(v)
     n = nu + nv - 1
     np2 = nextpow2(n)
     pad!(u,np2),pad!(v,np2)
     y = ifft_pow2(fft_pow2(u).*fft_pow2(v))
     #TODO This would not handle Dual/ComplexDual numbers correctly
     T = promote_type(mapreduce(eltype,promote_type,u),mapreduce(eltype,promote_type,v))
-    y = T<:Real ? real(y[1:n]) : y[1:n]
+    y = T<:Real ? real(@view y[1:n]) : y[1:n]
 end
 
 ######################################################################
@@ -33,9 +32,9 @@ end
 
 # plan_fft for BigFloats (covers Laurent svfft)
 
-plan_fft(x::Vector{F}) where {F<:Fun} = fft
-plan_ifft(x::Vector{F}) where {F<:Fun} = ifft
-plan_ifft!(x::Vector{F}) where {F<:Fun} = ifft
+plan_fft(x::Vector{<:Fun}) = fft
+plan_ifft(x::Vector{<:Fun}) = ifft
+plan_ifft!(x::Vector{<:Fun}) = ifft
 
 
 # Fourier space plans for BigFloat

test/NumberTypeTest.jl

@@ -1,4 +1,4 @@
-using ApproxFun, ApproxFunOrthogonalPolynomials, Test
+using ApproxFun, ApproxFunOrthogonalPolynomials, Test, FFTW
 
 @testset "BigFloat" begin
     @testset "BigFloat constructor" begin
@@ -80,4 +80,22 @@ using ApproxFun, ApproxFunOrthogonalPolynomials, Test
         w = (1-x^2)^b
         @test w(BigFloat(1)/10) ≈ (1-(BigFloat(1)/10)^2)^b
     end
+
+    @testset "fft" begin
+        @testset "Fun" begin
+            f = Fun(x->1, Fourier())
+            g = ApproxFun.fft(fill(f, 3))
+            @test g[1] ≈ Fun(x->3, Fourier())
+            @test g[2] ≈ Fun(x->0, Fourier()) atol=1e-10
+            @test g[3] ≈ Fun(x->0, Fourier()) atol=1e-10
+        end
+
+        @testset "BigFloat" begin
+            for n in 1:10
+                v = BigFloat[i for i in 1:n]
+                fv = ApproxFunBase.fft(v)
+                @test fv ≈ FFTW.fft(Float64.(v)) rtol=1e-8
+            end
+        end
+    end
 end