Merge pull request #30 from bionanoimaging/FourierJoin_fix

Fourier join fix

Author: roflmaostc

src/correlations.jl

@@ -1,6 +1,5 @@
 export ccorr
 
-
 """
     ccorr(u, v[, dims]; centered=false)
 
@@ -13,7 +12,6 @@ If either `u` or `v` is complex we use `fft` and output is hence complex.
 
 Per default the correlation is performed along `min(ndims(u), ndims(v))`.
 
-
 ```jldoctest
 julia> ccorr([1,1,0,0], [1,1,0,0], centered=true)
 4-element Vector{Float64}:

src/custom_fourier_types.jl

@@ -1,20 +1,31 @@
-##  This View checks for the index to be L1 or the mirrored version (L2)
-# and then replaces the value by half of the parent at L1
+"""
+    FourierSplit{T,N, AA<:AbstractArray{T, N}} <: AbstractArray{T,N}
+
+This View checks for the index to be L1 or the mirrored version (L2)
+and then replaces the value by half of the parent at L1
+`do_split` is a Bool that indicates whether this mechanism is active. 
+It is needed for type stability reasons of functions returnting this type.
+"""
 struct FourierSplit{T,N, AA<:AbstractArray{T, N}} <: AbstractArray{T,N}
     parent::AA # holds the data (or is another view)
     D::Int # dimension along which to apply to copy
     L1::Int # low index position to copy from (and half)
     L2::Int # high index positon to copy to (and half)
+    do_split::Bool
+
+"""
+    FourierSplit(parent::AA, D::Int,L1::Int,L2::Int, do_split::Bool) where {T,N, AA<:AbstractArray{T, N}}
 
-    # This version below is needed to avoid a split for the firs rft dimension but still return half the value
-    # FFTs and other RFT dimension should use the version without L2
-    function FourierSplit(parent::AA, D::Int,L1::Int,L2::Int) where {T,N, AA<:AbstractArray{T, N}}
-        return new{T,N, AA}(parent, D, L1, L2)
+This version below is needed to avoid a split for the first rft dimension,
+but still return half the value FFTs and other RFT dimension should use the version without L2.
+"""
+    function FourierSplit(parent::AA, D::Int,L1::Int,L2::Int, do_split::Bool) where {T,N, AA<:AbstractArray{T, N}}
+        return new{T,N, AA}(parent, D, L1, L2, do_split)
     end
-    function FourierSplit(parent::AA, D::Int,L1::Int) where {T,N, AA<:AbstractArray{T, N}}
+    function FourierSplit(parent::AA, D::Int, L1::Int, do_split::Bool) where {T,N, AA<:AbstractArray{T, N}}
         mid = fft_center(size(parent)[D])
         L2 = mid + (mid-L1)
-        return FourierSplit(parent, D,L1,L2)
+        return FourierSplit(parent, D,L1,L2, do_split)
     end
 end
 
@@ -25,7 +36,7 @@ Base.parent(A::FourierSplit) = A.parent
 Base.size(A::FourierSplit) = size(parent(A))
 
 @inline function Base.getindex(A::FourierSplit{T,N, <:AbstractArray{T, N}}, i::Vararg{Int,N}) where {T,N}
-    if i[A.D]==A.L2 || i[A.D]==A.L1 # index along this dimension A.D corrsponds to slice L2
+    if (i[A.D]==A.L2 || i[A.D]==A.L1) && A.do_split # index along this dimension A.D corrsponds to slice L2
         # not that "setindex" in the line below modifies only the index, not the array
         @inbounds return parent(A)[Base.setindex(i,A.L1, A.D)...] / 2
     else i[A.D]==A.L2
@@ -34,34 +45,45 @@ Base.size(A::FourierSplit) = size(parent(A))
     end
 end
 
-## This View checks for the index to be L1 
-# and then replaces the value by add the value at the mirrored position L2
+"""
+    FourierJoin{T,N, AA<:AbstractArray{T, N}} <: AbstractArray{T, N}
+
+This View checks for the index to be L1 
+and then replaces the value by add the value at the mirrored position L2
+`do_join` is a Bool that indicates whether this mechanism is active. 
+It is needed for type stability reasons of functions returnting this type
+"""
 struct FourierJoin{T,N, AA<:AbstractArray{T, N}} <: AbstractArray{T, N}
     parent::AA
     D::Int # dimension along which to apply to copy
     L1::Int # low index position to copy from (and half)
     L2::Int # high index positon to copy to (and half)
+    do_join::Bool
 
-    # This version below is needed to avoid a split for the firs rft dimension but still return half the value
-    # FFTs and other RFT dimension should use the version without L2
-    function FourierJoin(parent::AA, D::Int, L1::Int, L2::Int) where {T, N, AA<:AbstractArray{T, N}}
-        return new{T, N, AA}(parent, D, L1, L2)
+"""
+This version below is needed to avoid a split for the 
+first rft dimension but still return half the value
+FFTs and other RFT dimension should use the version without L2
+"""
+    function FourierJoin(parent::AA, D::Int, L1::Int, L2::Int, do_join::Bool) where {T, N, AA<:AbstractArray{T, N}}
+        return new{T, N, AA}(parent, D, L1, L2, do_join)
     end
 
-    function FourierJoin(parent::AA, D::Int,L1::Int) where {T, N, AA<:AbstractArray{T, N}}
+    function FourierJoin(parent::AA, D::Int,L1::Int, do_join::Bool) where {T, N, AA<:AbstractArray{T, N}}
         mid = fft_center(size(parent)[D])
         L2 = mid + (mid-L1)
-        return FourierJoin(parent, D, L1, L2)
+        return FourierJoin(parent, D, L1, L2, do_join)
     end
 end
+
 Base.IndexStyle(::Type{FS}) where {FS<:FourierJoin} = IndexStyle(parenttype(FS))
 parenttype(::Type{FourierJoin{T,N,AA}}) where {T,N,AA} = AA
 parenttype(A::FourierJoin) = parenttype(typeof(A))
 Base.parent(A::FourierJoin) = A.parent
 Base.size(A::FourierJoin) = size(parent(A))
 
 @inline function Base.getindex(A::FourierJoin{T,N, <:AbstractArray{T, N}}, i::Vararg{Int,N}) where {T,N}
-    if i[A.D]==A.L1
+    if i[A.D]==A.L1 && A.do_join
         @inbounds return (parent(A)[i...] + parent(A)[Base.setindex(i, A.L2, A.D)...])
     else 
         @inbounds return (parent(A)[i...])

src/czt.jl

@@ -78,6 +78,7 @@ end
 
 """
     czt(xin , scale, dims=1:length(size(xin)))
+
 Chirp z transform of the ND array `xin`
 This code is based on a 2D Matlab version of the CZT, written by H. Gross.
 The tuple `scale` defines the zoom factors in the Fourier domain. Each has to be bigger than one.
@@ -136,6 +137,7 @@ end
 
 """
     iczt(xin , scale, dims=1:length(size(xin)))
+
 Inverse chirp z transform of the ND array `xin`
 This code is based on a 2D Matlab version of the CZT, written by H. Gross.
 The tuple `scale` defines the zoom factors in the Fourier domain. Each has to be bigger than one.

src/fftshift_alternatives.jl

@@ -101,8 +101,8 @@ end
 """
     fftshift2d_view(mat::AbstractArray{T, N}) where {T, N}
 
-    Short-hand for `fftshift_view(mat, (1,2))`.
-    See `fftshift_view` for details.
+Short-hand for `fftshift_view(mat, (1,2))`.
+See `fftshift_view` for details.
 """
 function fftshift2d_view(mat::AbstractArray{T, N}) where {T, N}
     fftshift_view(mat,(1,2))
@@ -111,8 +111,8 @@ end
 """
     ifftshift2d_view(mat::AbstractArray{T, N}) where {T, N}
 
-    Short-hand for `ifftshift_view(mat, (1,2))` performing only a 2D inverse ft.
-    See `ifft` for details.
+Short-hand for `ifftshift_view(mat, (1,2))` performing only a 2D inverse ft.
+See `ifft` for details.
 """
 function ifftshift2d_view(mat::AbstractArray{T, N}) where {T, N}
     ifftshift_view(mat,(1,2))

src/fourier_filtering.jl

@@ -54,7 +54,7 @@ function fourier_filter!(arr::AbstractArray{<:Real, N}, wins::AbstractVector; tr
 end
 
 """
-fourier_filter(arr::AbstractArray, fct=window_gaussian; kwargs...)
+    fourier_filter(arr::AbstractArray, fct=window_gaussian; kwargs...)
     
 filters an array by multiplication in Fourierspace. This version uses in-place Fourier-transforms and multiplication whereever possible.
 The filter window function is assumed to be separable. Depending on the array type either full-complex or real-to-complex FFTs will be used.

src/fourier_resizing.jl

@@ -1,7 +1,6 @@
 export select_region_ft, select_region_rft
 export rft_fix_after, rft_fix_before, ft_fix_after, ft_fix_before
 
-
 """
     select_region_ft(mat,new_size)
 
@@ -62,12 +61,12 @@ julia> select_region_ft(ffts(x), (5,3))
  -86.75+0.0im     165.5+0.0im    -86.75+0.0im
 ```
 """
-function select_region_ft(mat,new_size)
+function select_region_ft(mat, new_size)
     old_size = size(mat)
-    mat_fixed_before = ft_fix_before(mat,old_size,new_size)
-    mat_pad = ft_pad(mat_fixed_before,new_size)
+    mat_fixed_before = ft_fix_before(mat, old_size, new_size)
+    mat_pad = ft_pad(mat_fixed_before, new_size)
     # afterwards we add the highest pos. frequency to the highest lowest one 
-    return ft_fix_after(mat_pad ,old_size,new_size)
+    return ft_fix_after(mat_pad ,old_size, new_size)
 end
 
 """
@@ -87,12 +86,11 @@ The size of the corresponding real-space array before it was rfted.
 The size of the corresponding real-space array view after the operation finished. The Fourier-center
 is always assumed to align before and after the padding aperation.
 """
-function select_region_rft(mat,old_size, new_size)
-    rft_old_size = size(mat)
-    rft_new_size = Base.setindex(new_size,new_size[1]÷2 +1, 1)
+function select_region_rft(mat, old_size, new_size)
+    # rft_old_size = size(mat)
+    rft_new_size = Base.setindex(new_size,new_size[1] ÷ 2 + 1, 1)
     return rft_fix_after(rft_pad(
-        rft_fix_before(mat,old_size,new_size),
-        rft_new_size),old_size,new_size)
+        rft_fix_before(mat, old_size, new_size), rft_new_size), old_size, new_size)
 end
 
 """
@@ -123,75 +121,112 @@ julia> select_region(ones(3,3),new_size=(7,7),center=(1,3))
 """
 function select_region(mat; new_size=size(mat), center=ft_center_diff(size(mat)).+1, pad_value=zero(eltype(mat)))
     new_size = Tuple(expand_size(new_size, size(mat)))
-    center = Tuple(expand_size(center, ft_center_diff(size(mat)).+1))
-    oldcenter = ft_center_diff(new_size).+1
-    PaddedView(pad_value, mat,new_size, oldcenter .- center.+1);
+    center = Tuple(expand_size(center, ft_center_diff(size(mat)) .+ 1))
+    oldcenter = ft_center_diff(new_size) .+ 1
+    PaddedView(pad_value, mat, new_size, oldcenter .- center.+1);
 end
 
 function ft_pad(mat, new_size)
-    return select_region(mat;new_size=new_size)
+    return select_region(mat; new_size = new_size)
 end
 
 function rft_pad(mat, new_size)
     c2 = rft_center_diff(size(mat))
-    c2 = Base.setindex(c2, new_size[1].÷2, 1);
-    return select_region(mat;new_size=new_size, center=c2.+1)
+    c2 = Base.setindex(c2, new_size[1] .÷ 2, 1);
+    return select_region(mat; new_size=new_size, center = c2 .+ 1)
 end
 
-function ft_fix_before(mat, size_old, size_new; start_dim=1)
-    for d = start_dim:ndims(mat)
-        sn = size_new[d]
-        so = size_old[d]
-        if sn < so && iseven(sn)
-            L1 = (size_old[d] -size_new[d] )÷2 +1
-            mat = FourierJoin(mat, d, L1)
-        end
-    end
-    return mat
+"""
+    ft_fix_before(mat::MT, size_old, size_new, ::Val{N})::FourierJoin{T,N,MT}  where {T,N, MT<:AbstractArray{T,N}} 
+
+implements the specialized (highest dimension) version of a recursive dimension-specific function that returns an array type which knows
+how to access (joins) certain elements.
+"""
+function ft_fix_before(mat::MT, size_old, size_new, ::Val{N})::FourierJoin{T,N,MT}  where {T,N, MT<:AbstractArray{T,N}} 
+    sn = size_new[N]
+    so = size_old[N]
+    do_join = (sn < so && iseven(sn))
+    L1 = (size_old[N] - size_new[N] ) ÷ 2 + 1
+    return FourierJoin(mat, N, L1, do_join)
 end
 
-function ft_fix_after(mat,size_old,size_new; start_dim=1)
-    start_dim
-    ndims(mat)
-    for d=start_dim:ndims(mat)
-        sn = size_new[d]
-        so = size_old[d]
-        if sn > so && iseven(so)
-            L1 = (size_new[d]-size_old[d])÷2+1
-            mat = FourierSplit(mat,d,L1)
-        end
-        # if equal do nothing
+"""
+    ft_fix_before(mat::MT, size_old, size_new, ::Val{D}=Val(1))  where {D, T, N, MT<:AbstractArray{T,N}} 
+
+implements the general version of a recursive dimension-specific function that returns an array type which knows
+how to access (joins) certain elements.
+"""
+function ft_fix_before(mat::MT, size_old, size_new, ::Val{D}=Val(1))  where {D, T, N, MT<:AbstractArray{T,N}} 
+    if D <= N
+        sn = size_new[D]
+        so = size_old[D]
+        do_join = (sn < so && iseven(sn))
+        L1 = (size_old[D] - size_new[D] )÷2 +1
+        mat = FourierJoin(mat, D, L1, do_join)
+        return ft_fix_before(mat, size_old, size_new, Val(D + 1))
+    else
+        L1 = (size_old[N] -size_new[N] )÷2 +1
+        return FourierJoin(mat, N, L1, false)
     end
-    return mat
 end
 
-function rft_fix_first_dim_before(mat,size_old,size_new;dim=1)
-    sn = size_new[dim] # Note that this dim is the corresponding real-space size
-    so = size_old[dim] # Note that this dim is the corresponding real-space size
-    if sn < so && iseven(sn) # result size is even upon cropping
-        L1 = size_new[dim] ÷ 2 + 1
-        mat = FourierJoin(mat, dim, L1, L1) # a hack to dublicate the value
+function ft_fix_after(mat::MT, size_old, size_new, ::Val{N})::FourierSplit{T,N,MT}   where {T, N, MT<:AbstractArray{T,N}}
+    sn = size_new[N]
+    so = size_old[N]
+    do_split = (sn > so && iseven(so))
+    L1 = (size_new[N] - size_old[N]) ÷ 2 + 1
+    return FourierSplit(mat, N, L1, do_split)
+end
+
+function ft_fix_after(mat::MT, size_old, size_new, ::Val{D}=Val(1)) where {D, T, N, MT<:AbstractArray{T,N}}
+    if D <= N
+        sn = size_new[D]
+        so = size_old[D]
+        do_split = (sn > so && iseven(so))
+        L1 = (size_new[D]-size_old[D])÷2+1
+        mat = FourierSplit(mat, D, L1, do_split)
+        return ft_fix_after(mat, size_old, size_new, Val(D + 1))
+    else
+        L1 = (size_new[N]-size_old[N])÷2+1
+        return FourierSplit(mat, N, L1, false)
     end
+end
+
+function rft_fix_first_dim_before(mat, size_old, size_new; dim=1)
+    # Note that this dim is the corresponding real-space size
+    sn = size_new[dim] 
+    so = size_old[dim]
+    # result size is even upon cropping
+    do_join = (sn < so && iseven(sn))
+    L1 = size_new[dim] ÷ 2 + 1
+    # a hack to dublicate the value
+    mat = FourierJoin(mat, dim, L1, L1, do_join)
     return mat
 end
 
 function rft_fix_first_dim_after(mat,size_old,size_new;dim=1)
-    sn = size_new[dim] # Note that this dim is the corresponding real-space size
-    so = size_old[dim] # Note that this dim is the corresponding real-space size
-    if sn > so && iseven(so) # source size is even upon padding
-        L1 = size_old[dim] ÷ 2 + 1
-        mat = FourierSplit(mat,dim,L1,-1) # This hack prevents a second position to be affected
-    end
+    # Note that this dim is the corresponding real-space size
+    sn = size_new[dim] 
+    so = size_old[dim] 
+    # source size is even upon padding
+    do_split = (sn > so && iseven(so)) 
+    L1 = size_old[dim] ÷ 2 + 1
+    # This hack prevents a second position to be affected
+    mat = FourierSplit(mat, dim, L1, -1, do_split)
     # if equal do nothing
     return mat
 end
 
 function rft_fix_before(mat,size_old,size_new)
-    mat=rft_fix_first_dim_before(mat,size_old,size_new;dim=1) # ignore the first dimension
-    ft_fix_before(mat,size_old,size_new;start_dim=2) # ignore the first dimension
+    # ignore the first dimension
+    mat=rft_fix_first_dim_before(mat,size_old,size_new;dim=1) 
+    # ignore the first dimension since it starts at Val(2)
+    ft_fix_before(mat, size_old, size_new, Val(2)) 
 end
 
 function rft_fix_after(mat,size_old,size_new)
-    mat = rft_fix_first_dim_after(mat,size_old,size_new;dim=1) # ignore the first dimension
-    ft_fix_after(mat,size_old,size_new;start_dim=2) # ignore the first dimension
+    # ignore the first dimension
+    mat = rft_fix_first_dim_after(mat,size_old,size_new;dim=1) 
+    # ignore the first dimension since it starts at Val(2)
+    ft_fix_after(mat, size_old, size_new, Val(2))
 end

src/nfft_nd.jl

@@ -169,6 +169,7 @@ end
 
 """
     plan_nfft_nd(src::AbstractArray{T,D}, dst_fkt::Function, dst_size=size(src); is_in_pixels=false, is_adjoint=false, kwargs...)
+
 Plans an n-dimensional non-uniform FFT on grids with a regular topology defined via the function `dst_fkt`.
 
 # Arguments

src/resampling.jl

@@ -270,6 +270,7 @@ end
 
 """
     barrel_pin(arr, rel=0.5)
+
 emulates a barrel (`rel>0`) or a pincushion (`rel<0`) distortion. The distortions are calculated using `resample_czt()` with separable quadratic zooms.
 
 See also: `resample_czt()`