allowed type system to infer array types and added filters

Author: RainerHeintzmann

Project.toml

@@ -11,6 +11,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NDTools = "98581153-e998-4eef-8d0d-5ec2c052313d"
 NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 PaddedViews = "5432bcbf-9aad-5242-b902-cca2824c8663"
+SeparableFunctions = "c8c7ead4-852c-491e-a42d-3d43bc74259e"
 ShiftedArrays = "1277b4bf-5013-50f5-be3d-901d8477a67a"
 
 [compat]

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

@@ -0,0 +1,83 @@
+export fourier_filter!, fourier_filter
+export filter_gaussian # , filter_gaussian2
+
+"""
+"""
+function fourier_filter!(arr::AbstractArray{<:Complex, N}, fct=window_gaussian; border_in=0.0, border_out=1.0) where {N}
+    return fourier_filter_by_1D_FT!(arr, fct; border_in=border_in, border_out=border_out)
+end
+
+function fourier_filter!(arr::AbstractArray{<:Int, N}, fct=window_gaussian; kwargs...) where {N}
+    throw(ArgumentError("FFTW.jl does not accept AbstractArrays{<:Int, N}. Convert array to a Real/Complex."))
+end
+
+function fourier_filter!(arr::AbstractArray{<:Real, N}, fct=window_gaussian; border_in=0.0, border_out=1.0) where {N}
+    return fourier_filter_by_1D_RFT!(arr, fct;  border_in=border_in, border_out=border_out)
+end
+
+"""
+    fourier_filter(arr, limits)
+
+Returning a fourier_filtered array.
+See [`fourier_filter!`](@ref fourier_filter!) for more details
+"""
+function fourier_filter(arr, fct=window_gaussian;  border_in=0.0, border_out=1.0)
+    return fourier_filter!(copy(arr), fct; border_in=border_in, border_out=border_out)
+end
+
+function fourier_filter_by_1D_FT!(arr::TA, fct=window_gaussian; border_in=0.0, border_out=1.0, dims=(1:ndims(arr))) where {N, TA <: AbstractArray{<:Complex, N}}
+    # iterates of the dimension d using the corresponding shift
+    sz = size(arr)
+    for d in dims #(d, limit) in pairs(limits)
+        w = TA(ifftshift(fct(real(eltype(arr)), select_sizes(arr,d), border_in=border_in, border_out=border_out)))
+        #w = TA(collect(ifftshift(fct(real(eltype(arr)), select_sizes(arr,d), border_in=border_in, border_out=border_out))))
+        # go to fourier space and apply w
+        fft!(arr, d)
+        arr .*= w
+    end
+
+    ifft!(arr, dims)
+    return arr
+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 fourier_filter_by_1D_RFT!(arr::TA, fct=window_gaussian;  border_in=0.0, border_out=1.0) where {T<:Real, N, TA<:AbstractArray{T, N}}
+    sz = size(arr)
+    for d = 1:length(sz) #(d, limit) in pairs(limits)
+        p = plan_rfft(arr, d)
+
+        arr_ft = p * arr
+        w = TA(fct(real(eltype(arr)), select_sizes(arr_ft,d), offset=CtrRFFT, scale=2 ./size(arr,d), border_in=border_in, border_out=border_out))
+        # w = TA(collect(fct(real(eltype(arr)), select_sizes(arr_ft,d), offset=CtrRFFT, scale=2 ./size(arr,d), border_in=border_in, border_out=border_out)))
+        arr_ft .*= w
+        # since we now did a single rfft dim, we can switch to the complex routine
+        # new_limits = ntuple(i -> i ≤ d ? 0 : limits[i], N)
+         # workaround to mimic in-place rfft
+        fourier_filter_by_1D_FT!(arr_ft;  border_in=border_in, border_out=border_out, dims=(2:ndims(arr_ft)))
+        # go back to real space now and return because shift_by_1D_FT processed
+        # the other dimensions already
+        mul!(arr, inv(p), arr_ft)
+        return arr # this breaks the for loop and finishes the algorithm
+    end
+    return arr
+end
+
+function filter_gaussian(arr, sigma=eltype(arr)(1))
+    sigma_k = 2 ./ (pi .* sigma) 
+    fourier_filter(arr, border_out=sigma_k)
+end
+
+# Suggested code by Felix. But it is not as fast as the code above:
+#
+# function filter_gaussian2(img::AbstractArray{T};
+#     sigma=one(T), dims=1:ndims(img)) where (T <: Real)
+# shiftdims = dims[2:end]
+# f = rfft(img, dims)
+# return irfft(f .* ifftshift(gaussian(T, size(f), 
+#                            offset=CtrRFT, 
+#                            sigma=size(img) ./ (T(2π) .*sigma)),
+#             shiftdims), 
+# size(img, dims[1]), dims)
+# end

src/fourier_shear.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,13 @@ 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 = TR(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,12 +59,14 @@ 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 = TR(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
@@ -103,10 +108,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))
-
+    # shift_strength = reshape(fftpos(1, size(arr, shear_dim), CenterFT), NDTools.select_sizes(arr, shear_dim))
+    # TR = real_arr_type(typeof(collect(arr[1:1]))) # There is a problem with circshifted arrays and this way of finding the type.
+    shift_strength = 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_shifting.jl

@@ -80,14 +80,17 @@ 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 = TR(reorient(fftfreq(size(arr, d)),d, Val(N)))
+        # freqs = TR(reshape(fftfreq(size(arr, d)), ntuple(i -> 1, Val(d-1))..., size(arr,d)))
+        # @show size(freqs)
         # allocates a 1D slice of exp values 
         if iszero(soft_fraction)
             ϕ = cispi.(- freqs .* 2 .* shift)
@@ -126,14 +129,21 @@ 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 = TR(reorient(fftfreq(size(arr, d))[1:s],d, Val(N)))
         if iszero(soft_fraction)
             ϕ = cispi.(-freqs .* 2 .* shift)
         else
@@ -146,9 +156,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)