Compute IFFT plan only once

JuliaAstro · Dec 16, 2016 · 9a450c6 · 9a450c6
1 parent 2e07d53
commit 9a450c6
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 16 deletions.
diff --git a/src/LombScargle.jl b/src/LombScargle.jl
@@ -202,11 +202,18 @@ function _press_rybicki{R1<:Real,R2<:Real,R3<:Real,R4<:Real}(times::AbstractVect
     N  = length(freqs)
     f0 = minimum(freqs)
     #---------------------------------------------------------------------------
+    # 0. prepare for the computation of IFFT.
+    # `ifft' can take arrays of any length in input, but
+    # it's faster when the length is exactly a power of 2.
+    Nfft = nextpow2(N * oversampling)
+    plan = plan_ifft(Vector{Complex{promote_type(R3, R2)}}(Nfft),
+                     flags = FFTW.MEASURE)
+    #---------------------------------------------------------------------------
     # 1. compute functions of the time-shift tau at each frequency
-    Ch, Sh = trig_sum(times, w .* y, df, N, f0, 1, oversampling, Mfft)
-    C2, S2 = trig_sum(times, w,      df, N, f0, 2, oversampling, Mfft)
+    Ch, Sh = trig_sum(plan, times, w .* y, df, N, Nfft, f0, 1, Mfft)
+    C2, S2 = trig_sum(plan, times, w,      df, N, Nfft, f0, 2, Mfft)
     if fit_mean
-        C, S = trig_sum(times, w, df, N, f0, 1, oversampling, Mfft)
+        C, S = trig_sum(plan, times, w, df,    N, Nfft, f0, 1, Mfft)
         tan_2ωτ = (S2 .- 2.0 * S .* C) ./ (C2 .- (C .* C .- S .* S))
     else
         tan_2ωτ = S2 ./ C2

diff --git a/src/extirpolation.jl b/src/extirpolation.jl
@@ -20,9 +20,9 @@
 #
 ### Code:
 
-function add_at!{N1<:Number,N2<:Number,N3<:Number}(arr::AbstractVector{N1},
-                                                   ind::AbstractVector{N2},
-                                                   vals::AbstractVector{N3})
+function add_at!{T1,I<:Integer,T2}(arr::AbstractVector{T1},
+                                   ind::AbstractVector{I},
+                                   vals::AbstractVector{T2})
     @inbounds for i in eachindex(ind)
         arr[ind[i]] += vals[i]
     end
@@ -57,25 +57,26 @@ function extirpolate{RE<:Real,NU<:Number}(X::AbstractVector{RE},
     return result
 end
 
-function trig_sum{R1<:Real,R2<:Real}(t::AbstractVector{R1},
+function trig_sum{R1<:Real,R2<:Real}(ifft_plan,
+                                     t::AbstractVector{R1},
                                      h::AbstractVector{R2}, df::Real,
-                                     N::Integer, f0::Real=0.0,
+                                     N::Integer, Nfft::Integer,
+                                     f0::Real=0.0,
                                      freq_factor::Integer=1,
-                                     oversampling::Integer=5, Mfft::Integer=4)
+                                     Mfft::Integer=4)
     @assert Mfft > 0
     df *= freq_factor
     f0 *= freq_factor
     @assert df > 0
-    # `ifft' can take arrays of any length in input, but it's faster when the
-    # length is exactly a power of 2.
-    Nfft = nextpow2(N * oversampling)
     t0 = minimum(t)
     if f0 > 0
-        h = h .* exp(2im * pi * f0 * (t - t0))
+        H = h .* exp(2im * pi * f0 * (t - t0))
+    else
+        H = complex(h)
     end
     tnorm = mod(((t - t0) * Nfft * df), Nfft)
-    grid = extirpolate(tnorm, h, Nfft, Mfft)
-    fftgrid = Nfft * ifft(grid)[1:N]
+    grid = extirpolate(tnorm, H, Nfft, Mfft)
+    fftgrid = Nfft * (ifft_plan * grid)[1:N]
     if t0 != 0
         f = f0 + df * (0:N - 1)
         fftgrid .*= exp(2im * pi * t0 * f)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -133,7 +133,10 @@ y = sin(x)
 @test_approx_eq LombScargle.extirpolate(x, y, 11) LombScargle.extirpolate(x, y, 13)[1:11]
 
 # Test trig_sum
-C, S = LombScargle.trig_sum(x, y, 1, 10)
+N = 10
+Nfft = nextpow2(5N)
+p = plan_ifft(Vector{Complex{Float64}}(Nfft))
+C, S = LombScargle.trig_sum(p, x, y, 1, N, Nfft)
 @test_approx_eq S [0.0,0.3753570125888358,0.08163980192703546,-0.10139634351774979,-0.4334223744905633,-2.7843373311769777,0.32405810159838055,0.05729663600471602,-0.13191736591325876,-0.5295781583202946]
 @test_approx_eq C [8.708141477890015,-0.5402668064176129,-0.37460815054027985,-0.3793457539084364,-0.5972351546196192,14.612204307982497,-0.5020253140297526,-0.37724493022381034,-0.394096831764578,-0.6828241623474718]