Simplify Kernels.init_fourier_coefficients!

jipolanco · Oct 15, 2024 · ab21f69 · ab21f69
1 parent d40402b
commit ab21f69
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 17 deletions.
diff --git a/src/Kernels/Kernels.jl b/src/Kernels/Kernels.jl
@@ -64,7 +64,8 @@ function init_fourier_coefficients!(g::AbstractKernelData, ks::AbstractVector)
     Nk == length(gk) && return gk  # assume coefficients were already computed
     resize!(gk, Nk)
     @assert eachindex(gk) == eachindex(ks)
-    evaluate_fourier!(g, gk, ks)
+    f = evaluate_fourier_func(g)
+    map!(f, gk, ks)
     gk
 end
 

diff --git a/src/Kernels/bspline.jl b/src/Kernels/bspline.jl
@@ -95,11 +95,9 @@ function evaluate_kernel_func(g::BSplineKernelData{M}) where {M}
     end
 end
 
-# This should work on CPU and GPU.
-function evaluate_fourier!(g::BSplineKernelData{M}, gk::AbstractVector, ks::AbstractVector) where {M}
+function evaluate_fourier_func(g::BSplineKernelData{M}) where {M}
     (; Δt,) = g
-    @assert eachindex(gk) == eachindex(ks)
-    map!(gk, ks) do k
+    function (k)
         kh = k * Δt / 2
         s = sin(kh) / kh
         n = 2M

diff --git a/src/Kernels/gaussian.jl b/src/Kernels/gaussian.jl
@@ -98,11 +98,9 @@ function optimal_kernel(kernel::GaussianKernel, h::HalfSupport{M}, Δx, σ; back
     GaussianKernelData(h, backend, Δx, ℓ)
 end
 
-# This should work on CPU and GPU.
-function evaluate_fourier!(g::GaussianKernelData, gk::AbstractVector, ks::AbstractVector)
+function evaluate_fourier_func(g::GaussianKernelData)
     (; τ,) = g
-    @assert eachindex(gk) == eachindex(ks)
-    map!(gk, ks) do k
+    function (k)
         exp(-τ * k^2 / 4) * sqrt(π * τ)  # = exp(-σ² k² / 2) * sqrt(2πσ²)
     end
 end

diff --git a/src/Kernels/kaiser_bessel.jl b/src/Kernels/kaiser_bessel.jl
@@ -141,11 +141,9 @@ function optimal_kernel(kernel::KaiserBesselKernel, h::HalfSupport{M}, Δx, σ;
     KaiserBesselKernelData(h, backend, Δx, β)
 end
 
-# This should work on CPU and GPU.
-function evaluate_fourier!(g::KaiserBesselKernelData, gk::AbstractVector, ks::AbstractVector)
-    @assert eachindex(gk) == eachindex(ks)
+function evaluate_fourier_func(g::KaiserBesselKernelData)
     (; β², w,) = g
-    map!(gk, ks) do k
+    function (k)
         q = w * k
         s = sqrt(β² - q^2)  # this is always real (assuming β ≥ Mπ)
         2 * w * sinh(s) / s

diff --git a/src/Kernels/kaiser_bessel_backwards.jl b/src/Kernels/kaiser_bessel_backwards.jl
@@ -112,11 +112,9 @@ function optimal_kernel(kernel::BackwardsKaiserBesselKernel, h::HalfSupport{M},
     BackwardsKaiserBesselKernelData(h, backend, Δx, β)
 end
 
-# This should work on CPU and GPU.
-function evaluate_fourier!(g::BackwardsKaiserBesselKernelData, gk::AbstractVector, ks::AbstractVector)
-    @assert eachindex(gk) == eachindex(ks)
+function evaluate_fourier_func(g::BackwardsKaiserBesselKernelData)
     (; β, w,) = g
-    map!(gk, ks) do k
+    function (k)
         q = w * k
         s = sqrt(β^2 - q^2)  # this is always real (assuming β ≥ Mπ)
         w * besseli0(s)