Merge pull request #23 from bionanoimaging/fourier_filtering

Fourier filtering

Author: roflmaostc

Project.toml

@@ -18,7 +18,7 @@ ChainRulesCore = "1, 1.0, 1.1"
 FFTW = "1.5"
 ImageTransformations = "0.9"
 IndexFunArrays = "0.2"
-NDTools = "0.5"
+NDTools = "0.5.1"
 NFFT = "0.11, 0.12, 0.13"
 PaddedViews = "0.5"
 ShiftedArrays = "2"

src/FourierTools.jl

@@ -15,6 +15,7 @@ include("resampling.jl")
 include("custom_fourier_types.jl")
 include("fourier_resizing.jl")
 include("fourier_shifting.jl")
+include("fourier_filtering.jl")
 include("fourier_resample_1D_based.jl")
 include("fourier_rotate.jl")
 include("fourier_shear.jl")

src/fourier_filtering.jl

@@ -16,16 +16,17 @@ There is also `shear!` available.
 + `fix_nyquist`: apply a fix to the highest frequency during the Fourier-space application of the exponential factor
 + `adapt_size`: if true, pad the data prior to the shear. The result array will be larger
 + `pad_value`: the value to pad with (only applies if `adapt_size=true`)
++ `assign_wrap=assign_wrap`: replaces wrap-around areas by `pad_value` (only of `adapt_size` is `false`)
 
 For complex arrays we use `fft`, for real array we use `rfft`.
 """
-function shear(arr::AbstractArray, Δ, shear_dir_dim=1, shear_dim=2; fix_nyquist=false, adapt_size=false::Bool, pad_value=zero(eltype(arr)))
+function shear(arr::AbstractArray, Δ, shear_dir_dim=1, shear_dim=2; fix_nyquist=false, assign_wrap=false, adapt_size=false::Bool, pad_value=zero(eltype(arr)))
     if adapt_size
         ns = Tuple(d == shear_dir_dim ? size(arr,d)+ceil(Int,abs.(Δ)) : size(arr,d) for d in 1:ndims(arr))
         arr2 = collect(select_region(arr, new_size=ns, pad_value=pad_value))
         return shear!(arr2, Δ, shear_dir_dim, shear_dim, fix_nyquist=fix_nyquist)
     else
-        return shear!(copy(arr), Δ, shear_dir_dim, shear_dim, fix_nyquist=fix_nyquist)
+        return shear!(copy(arr), Δ, shear_dir_dim, shear_dim, fix_nyquist=fix_nyquist, assign_wrap=assign_wrap, pad_value=pad_value)
     end
 end
 
@@ -40,11 +41,14 @@ For more details see `shear.`
 For complex arrays we can completely avoid large memory allocations.
 For real arrays, we need at least allocate on array in the fourier space.
 """
-function shear!(arr::AbstractArray{<:Complex, N}, Δ, shear_dir_dim=1, shear_dim=2; fix_nyquist=false, assign_wrap=false, pad_value=zero(eltype(arr))) where N
+function shear!(arr::TA, Δ, shear_dir_dim=1, shear_dim=2; fix_nyquist=false, assign_wrap=false, pad_value=zero(eltype(arr))) where {N, TA<:AbstractArray{<:Complex, N}}
     fft!(arr, shear_dir_dim)
 
     # stores the maximum amount of shift
-    shift = reshape(fftfreq(size(arr, shear_dir_dim)), NDTools.select_sizes(arr, shear_dir_dim))
+    # TR = real_arr_type(TA)
+    shift = similar(arr, real(eltype(arr)), select_sizes(arr, shear_dir_dim))
+    shift .= reshape(fftfreq(size(arr, shear_dir_dim)), NDTools.select_sizes(arr, shear_dir_dim))
+    # shift = TR(reorient(fftfreq(size(arr, shear_dir_dim)), shear_dir_dim, Val(N)))
 
     apply_shift_strength!(arr, arr, shift, shear_dir_dim, shear_dim, Δ, fix_nyquist)
 
@@ -56,16 +60,19 @@ function shear!(arr::AbstractArray{<:Complex, N}, Δ, shear_dir_dim=1, shear_dim
     return arr
 end
 
-function shear!(arr::AbstractArray{<:Real, N}, Δ, shear_dir_dim=1, shear_dim=2; fix_nyquist=false, assign_wrap=false, pad_value=zero(eltype(arr))) where N
+function shear!(arr::TA, Δ, shear_dir_dim=1, shear_dim=2; fix_nyquist=false, assign_wrap=false, pad_value=zero(eltype(arr))) where {N, TA<:AbstractArray{<:Real, N}}
     p = plan_rfft(arr, shear_dir_dim)
     arr_ft = p * arr 
 
     # stores the maximum amount of shift
-    shift = reshape(rfftfreq(size(arr, shear_dir_dim)), NDTools.select_sizes(arr_ft, shear_dir_dim))
+    # TR = real_arr_type(TA)
+    shift = similar(arr, real(eltype(arr_ft)), select_sizes(arr_ft, shear_dir_dim))
+    shift .= reshape(rfftfreq(size(arr, shear_dir_dim)), NDTools.select_sizes(arr_ft, shear_dir_dim))
+    # shift = TR(reorient(rfftfreq(size(arr, shear_dir_dim)),shear_dir_dim, Val(N)))
 
     apply_shift_strength!(arr_ft, arr, shift, shear_dir_dim, shear_dim, Δ, fix_nyquist)
     # go back to real space
-     
+
     # overwrites arr in-place
     ldiv!(arr, p, arr_ft)
     if assign_wrap
@@ -103,10 +110,12 @@ function assign_shear_wrap!(arr, Δ, shear_dir_dim=1, shear_dim=2, pad_value=zer
     end
 end
 
-function apply_shift_strength!(arr, arr_orig, shift, shear_dir_dim, shear_dim, Δ, fix_nyquist=false)
+function apply_shift_strength!(arr::TA, arr_orig, shift, shear_dir_dim, shear_dim, Δ, fix_nyquist=false) where {T, N, TA<:AbstractArray{T, N}}
     #applies the strength to each slice
-    shift_strength = reshape(fftpos(1, size(arr, shear_dim), CenterFT), NDTools.select_sizes(arr, shear_dim))
-
+    # The TR trick does not seem to work for the code below due to a call with a PaddedArray.
+    shift_strength = similar(arr, real(eltype(arr)), select_sizes(arr, shear_dim))
+    shift_strength .= (real(eltype(TA))).(reorient(fftpos(1, size(arr, shear_dim), CenterFT), shear_dim, Val(N)))
+ 
     # do the exp multiplication in place
     e = cispi.(2 .* Δ .* shift .* shift_strength)
     # for even arrays we need to fix real property of highest frequency

src/fourier_shear.jl

@@ -80,20 +80,26 @@ function soft_shift(freqs, shift, fraction=eltype(freqs)(0.1); corner=false)
     return cispi.(-freqs .* 2 .* (w .* shift + (1.0 .-w).* rounded_shift))
 end
 
-function shift_by_1D_FT!(arr::AbstractArray{<:Complex, N}, shifts; soft_fraction=0, take_real=false, fix_nyquist_frequency=false) where {N}
+function shift_by_1D_FT!(arr::TA, shifts; soft_fraction=0, take_real=false, fix_nyquist_frequency=false) where {N, TA<:AbstractArray{<:Complex, N}}
     # iterates of the dimension d using the corresponding shift
     for (d, shift) in pairs(shifts)
         if iszero(shift)
             continue
         end
         # better use reorient from NDTools here?
-        freqs = reshape(fftfreq(size(arr, d)), ntuple(i -> 1, Val(d-1))..., size(arr,d))
+        # TR = real_arr_type(TA)
+
+        freqs = similar(arr, real(eltype(arr)), select_sizes(arr, d))
+        # freqs = TR(reorient(fftfreq(size(arr, d)),d, Val(N)))
+        freqs .= reorient(fftfreq(size(arr, d)),d, Val(N))
+        # @show size(freqs)
         # allocates a 1D slice of exp values 
         if iszero(soft_fraction)
             ϕ = cispi.(- freqs .* 2 .* shift)
         else
             ϕ = soft_shift(freqs, shift, soft_fraction)
         end
+        # ϕ = exp_ikx_sep(complex_arr_type(TA), size(arr), dims=(d,), shift_by = shift)[1]
         # in even case, set one value to real
         if iseven(size(arr, d))
             s = size(arr, d) ÷ 2 + 1
@@ -126,14 +132,24 @@ end
 # the idea is the following:
 # rfft(x, 1) -> exp shift -> fft(x, 2) -> exp shift ->  fft(x, 3) -> exp shift -> ifft(x, [2,3]) -> irfft(x, 1)
 # So once we did a rft to shift something we can call the routine for complex arrays to shift
-function shift_by_1D_RFT!(arr::AbstractArray{<:Real, N}, shifts; soft_fraction=0, fix_nyquist_frequency=false, take_real=true) where {T, N}
+function shift_by_1D_RFT!(arr::TA, shifts; soft_fraction=0, fix_nyquist_frequency=false, take_real=true) where {N, TA<:AbstractArray{<:Real, N}}
     for (d, shift) in pairs(shifts)
         if iszero(shift)
             continue
         end
 
-        s = size(arr, d) ÷ 2 + 1
-        freqs = reshape(fftfreq(size(arr, d))[1:s], ntuple(i -> 1, d-1)..., s) 
+        p = plan_rfft(arr, d)
+        arr_ft = p * arr
+
+        #s1 = select_sizes(arr_ft, d)
+        s = size(arr_ft,d); # s1[d]
+        # @show size(arr, d) ÷ 2 + 1
+        # freqs = TR(reshape(fftfreq(size(arr, d))[1:s], ntuple(i -> 1, d-1)..., s1))
+
+        # TR = real_arr_type(TA)
+        freqs = similar(arr, real(eltype(arr_ft)), select_sizes(arr_ft, d))
+        # freqs = TR(reorient(fftfreq(size(arr, d))[1:s], d, Val(N)))
+        freqs .= reorient(rfftfreq(size(arr, d)), d, Val(N))
         if iszero(soft_fraction)
             ϕ = cispi.(-freqs .* 2 .* shift)
         else
@@ -146,9 +162,6 @@ function shift_by_1D_RFT!(arr::AbstractArray{<:Real, N}, shifts; soft_fraction=0
             invr = isinf(invr) ? 0 : invr
             ϕ[s] = fix_nyquist_frequency ? invr : ϕ[s]
         end
-        p = plan_rfft(arr, d)
-
-        arr_ft = p * arr
         arr_ft .*= ϕ
         # since we now did a single rfft dim, we can switch to the complex routine
         new_shifts = ntuple(i -> i ≤ d ? 0 : shifts[i], N)

src/fourier_shifting.jl

@@ -0,0 +1,36 @@
+Random.seed!(42)
+
+@testset "Fourier filtering" begin
+
+    @testset "Gaussian filter complex" begin
+        sz = (21, 22)
+        x = randn(ComplexF32, sz)
+        sigma = (1.1,2.2)
+        gf = filter_gaussian(x, sigma, real_space_kernel=false)
+        # Note that this is not the same, since one kernel is generated in real space and one in Fourier space!
+        # with sizes around 10, the difference is huge!
+        k = gaussian(Float32, sz, sigma=sigma)
+        k = k./sum(k) # different than "normal".
+        gfc = conv_psf(x, k)
+        @test ≈(gf,gfc, rtol=1e-2) # it is realatively inaccurate due to the kernel being generated in different places
+        gfr = filter_gaussian(x, sigma, real_space_kernel=true)
+        @test ≈(gfr,gfc) # it can be debated how to best normalize a Gaussian filter
+    end
+
+    @testset "Gaussian filter real" begin
+        sz = (21, 22)
+        x = randn(Float32, sz)
+        sigma = (1.1,2.2)
+        gf = filter_gaussian(x, sigma, real_space_kernel=true)
+        # Note that this is not the same, since one kernel is generated in real space and one in Fourier space!
+        # with sizes around 10, the difference is huge!
+        k = gaussian(sz, sigma=sigma)
+        k = k./sum(k) # different than "normal".
+        gf2 = conv_psf(x, k)
+        @test ≈(gf,gf2, rtol=1e-2) # it is realatively inaccurate due to the kernel being generated in different places
+    end
+    @testset "Other filters" begin
+        @test filter_hamming(FourierTools.delta(Float32, (3,)), border_in=0.0, border_out=1.0) ≈ [0.23,0.54, 0.23]
+        @test filter_hann(FourierTools.delta(Float32, (3,)), border_in=0.0, border_out=1.0) ≈ [0.25,0.5, 0.25]
+    end
+end

test/fourier_filtering.jl

@@ -21,5 +21,6 @@ include("custom_fourier_types.jl")
 include("damping.jl")
 include("czt.jl")
 include("nfft_tests.jl")
+include("fourier_filtering.jl")
 
 return