Add trig_sum function

JuliaAstro · Aug 17, 2016 · dbf1cc7 · dbf1cc7
1 parent 633627e
commit dbf1cc7
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/src/extirpolation.jl b/src/extirpolation.jl
@@ -30,6 +30,7 @@ end
 function extirpolate{R1<:Real,R2<:Real}(X::AbstractVector{R1},
                                         Y::AbstractVector{R2},
                                         N::Integer=0, M::Integer=4)
+    @assert length(X) == length(Y)
     x, y = collect(X), collect(Y)
     if N <= 0
         # Get the maximum of "X", `maximum' has a faster method for ranges.
@@ -43,7 +44,7 @@ function extirpolate{R1<:Real,R2<:Real}(X::AbstractVector{R1},
     ilo = clamp(trunc(Int, x - div(M, 2)), 0, N - M)
     numerator = y .* [prod(x[j] - ilo[j] - (0:M-1)) for j in eachindex(x)]
     denominator = factorial(M - 1)
-    for j in 0:(M - 1)
+    @inbounds for j in 0:(M - 1)
         if j > 0
             denominator *= j/(j - M)
         end
@@ -52,3 +53,29 @@ function extirpolate{R1<:Real,R2<:Real}(X::AbstractVector{R1},
     end
     return result
 end
+
+function trig_sum{R1<:Real,R2<:Real}(t::AbstractVector{R1},
+                                     h::AbstractVector{R2}, df::Real,
+                                     N::Integer, f0::Real=0.0,
+                                     freq_factor::Integer=1,
+                                     oversampling::Integer=5, Mfft::Integer=4)
+    @assert Mfft > 0
+    df *= freq_factor
+    f0 *= freq_factor
+    @assert df > 0
+    Nfft = nextpow2(N * oversampling)
+    t0 = minimum(t)
+    if f0 > 0
+        h = h * exp(2im * pi * f0 * (t - t0))
+    end
+    tnorm = mod(((t - t0) * Nfft * df), Nfft)
+    grid = extirpolate(tnorm, h, Nfft, Mfft)
+    fftgrid = Nfft * ifft(grid)[1:N]
+    if t0 != 0
+        f = f0 + df * 0:(N - 1)
+        fftgrid *= exp(2im * pi * t0 * f)
+    end
+    S = real(fftgrid)
+    C = imag(fftgrid)
+    return C, S
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -79,3 +79,8 @@ x = linspace(0, 10)
 y = sin(x)
 @test_approx_eq LombScargle.extirpolate(x, y) [0.39537718210649553,3.979484140636793,4.833090108345013,0.506805556164743,-3.828112427525919,-4.748341359084166,-1.3022050566901917,3.3367666084342256,5.070478111668922,1.291245296032218,-0.8264466821981216,0.0,0.0]
 @test_approx_eq LombScargle.extirpolate(x, y, 11) LombScargle.extirpolate(x, y)[1:11]
+
+# Test trig_sum
+C, S = LombScargle.trig_sum(x, y, 1, 10)
+@test_approx_eq C [0.0,0.3753570125888358,0.08163980192703546,-0.10139634351774979,-0.4334223744905633,-2.7843373311769777,0.32405810159838055,0.05729663600471602,-0.13191736591325876,-0.5295781583202946]
+@test_approx_eq S [8.708141477890015,-0.5402668064176129,-0.37460815054027985,-0.3793457539084364,-0.5972351546196192,14.612204307982497,-0.5020253140297526,-0.37724493022381034,-0.394096831764578,-0.6828241623474718]