From a7ebdac5bd9bb1f362862dcbd5210652ea1987ac Mon Sep 17 00:00:00 2001 From: Grigorii Starkov Date: Wed, 16 Mar 2022 19:08:57 +0100 Subject: [PATCH 1/4] Adds nonallocating versions of `fftshift` and `ifftshift` with the corresponding tests. Fixes #62. --- src/definitions.jl | 22 ++++++++++++++++++++++ test/runtests.jl | 17 +++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/src/definitions.jl b/src/definitions.jl index 80a6656..3dd40a0 100644 --- a/src/definitions.jl +++ b/src/definitions.jl @@ -364,6 +364,17 @@ function fftshift(x, dim = 1:ndims(x)) circshift(x, s) end +""" + fftshift!(dest, src, [dim]) + +Nonallocating version of [`fftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array. +""" +function fftshift!(dest, src, dim = 1:ndims(src)) + @assert size(dest)==size(src) + s = ntuple(d -> d in dim ? div(size(dest,d),2) : 0, Val(ndims(dest))) + circshift!(dest, src, s) +end + """ ifftshift(x, [dim]) @@ -384,6 +395,17 @@ function ifftshift(x, dim = 1:ndims(x)) circshift(x, s) end +""" + ifftshift!(dest, src, [dim]) + +Nonallocating version of [`ifftshift`](@ref). Stores the result of the shift of the `src` array into the `dest array`. +""" +function ifftshift!(dest, src, dim = 1:ndims(src)) + @assert size(dest)==size(src) + s = ntuple(d -> d in dim ? -div(size(src,d),2) : 0, Val(ndims(src))) + circshift!(dest, src, s) +end + ############################################################################## diff --git a/test/runtests.jl b/test/runtests.jl index de0304d..869c745 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -103,20 +103,37 @@ end @test @inferred(AbstractFFTs.fftshift([1 2 3])) == [3 1 2] @test @inferred(AbstractFFTs.fftshift([1, 2, 3])) == [3, 1, 2] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6])) == [6 4 5; 3 1 2] + a = [0 0 0] + b = [0, 0, 0] + c = [0 0 0; 0 0 0] + @test (AbstractFFTs.fftshift!(a, [1 2 3]); a == [3 1 2]) + @test (AbstractFFTs.fftshift!(b, [1, 2, 3]); b == [3, 1, 2]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6]); c == [6 4 5; 3 1 2]) @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], 1)) == [4 5 6; 1 2 3] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], ())) == [1 2 3; 4 5 6] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], (1,2))) == [6 4 5; 3 1 2] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], 1:2)) == [6 4 5; 3 1 2] + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [6 4 5; 3 1 2]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1:2); c == [6 4 5; 3 1 2]) @test @inferred(AbstractFFTs.ifftshift([1 2 3])) == [2 3 1] @test @inferred(AbstractFFTs.ifftshift([1, 2, 3])) == [2, 3, 1] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6])) == [5 6 4; 2 3 1] + @test (AbstractFFTs.ifftshift!(a, [1 2 3]); a == [2 3 1]) + @test (AbstractFFTs.ifftshift!(b, [1, 2, 3]); b == [2, 3, 1]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6]); c == [5 6 4; 2 3 1]) @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], 1)) == [4 5 6; 1 2 3] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], ())) == [1 2 3; 4 5 6] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], (1,2))) == [5 6 4; 2 3 1] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], 1:2)) == [5 6 4; 2 3 1] + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [5 6 4; 2 3 1]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1:2); c == [5 6 4; 2 3 1]) end @testset "FFT Frequencies" begin From 8c9d817c08f868aed57f88a15850f8fc8f42bd43 Mon Sep 17 00:00:00 2001 From: Grigorii Starkov Date: Wed, 16 Mar 2022 20:33:00 +0100 Subject: [PATCH 2/4] Adds `fftshift!` and `ifftshift!` to exported functions. --- src/AbstractFFTs.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractFFTs.jl b/src/AbstractFFTs.jl index 7b0902f..734c7d4 100644 --- a/src/AbstractFFTs.jl +++ b/src/AbstractFFTs.jl @@ -5,7 +5,7 @@ import ChainRulesCore export fft, ifft, bfft, fft!, ifft!, bfft!, plan_fft, plan_ifft, plan_bfft, plan_fft!, plan_ifft!, plan_bfft!, rfft, irfft, brfft, plan_rfft, plan_irfft, plan_brfft, - fftshift, ifftshift, Frequencies, fftfreq, rfftfreq + fftshift, ifftshift, fftshift!, ifftshift!, Frequencies, fftfreq, rfftfreq include("definitions.jl") include("chainrules.jl") From c562ae1efb7820607bd85f9988a9b29972731fcf Mon Sep 17 00:00:00 2001 From: Grigorii Starkov Date: Thu, 31 Mar 2022 13:16:37 +0200 Subject: [PATCH 3/4] Rewrote allocating shift functions to use non-allocating ones. Added backreferences to non-allocating functions. --- src/definitions.jl | 42 +++++++++++++++++++++++------------------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/src/definitions.jl b/src/definitions.jl index 3dd40a0..5fd7e1e 100644 --- a/src/definitions.jl +++ b/src/definitions.jl @@ -344,6 +344,17 @@ plan_irfft ############################################################################## +""" + fftshift!(dest, src, [dim]) + +Nonallocating version of [`fftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array. +""" +function fftshift!(dest, src, dim = 1:ndims(src)) + @assert size(dest)==size(src) + s = ntuple(d -> d in dim ? div(size(dest,d),2) : 0, Val(ndims(dest))) + circshift!(dest, src, s) +end + """ fftshift(x, [dim]) @@ -356,22 +367,24 @@ swapping the first and second halves, so `fftshift` and [`ifftshift`](@ref) are the same. If `dim` is not given then the signal is shifted along each dimension. + +The output of `fftshift` is allocated. If one desires to store the output in a preallocated array, use [`fftshift!`](@ref) instead. """ fftshift function fftshift(x, dim = 1:ndims(x)) - s = ntuple(d -> d in dim ? div(size(x,d),2) : 0, Val(ndims(x))) - circshift(x, s) + dest = similar(x) + fftshift!(dest, x, dim) end """ - fftshift!(dest, src, [dim]) + ifftshift!(dest, src, [dim]) -Nonallocating version of [`fftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array. +Nonallocating version of [`ifftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array. """ -function fftshift!(dest, src, dim = 1:ndims(src)) +function ifftshift!(dest, src, dim = 1:ndims(src)) @assert size(dest)==size(src) - s = ntuple(d -> d in dim ? div(size(dest,d),2) : 0, Val(ndims(dest))) + s = ntuple(d -> d in dim ? -div(size(src,d),2) : 0, Val(ndims(src))) circshift!(dest, src, s) end @@ -387,23 +400,14 @@ swapping the first and second halves, so [`fftshift`](@ref) and `ifftshift` are the same. If `dim` is not given then the signal is shifted along each dimension. + +The output of `ifftshift` is allocated. If one desires to store the output in a preallocated array, use [`ifftshift!`](@ref) instead. """ ifftshift function ifftshift(x, dim = 1:ndims(x)) - s = ntuple(d -> d in dim ? -div(size(x,d),2) : 0, Val(ndims(x))) - circshift(x, s) -end - -""" - ifftshift!(dest, src, [dim]) - -Nonallocating version of [`ifftshift`](@ref). Stores the result of the shift of the `src` array into the `dest array`. -""" -function ifftshift!(dest, src, dim = 1:ndims(src)) - @assert size(dest)==size(src) - s = ntuple(d -> d in dim ? -div(size(src,d),2) : 0, Val(ndims(src))) - circshift!(dest, src, s) + dest = similar(x) + ifftshift!(dest, x, dim) end ############################################################################## From bdcf67070db120dc686388d3c4489ad9b30ccb21 Mon Sep 17 00:00:00 2001 From: Grigorii Starkov Date: Thu, 31 Mar 2022 17:11:09 +0200 Subject: [PATCH 4/4] Fixed indenting, deleted asserts, made one-liners. --- src/definitions.jl | 16 ++++------------ test/runtests.jl | 28 ++++++++++++++-------------- 2 files changed, 18 insertions(+), 26 deletions(-) diff --git a/src/definitions.jl b/src/definitions.jl index 5fd7e1e..41df3e5 100644 --- a/src/definitions.jl +++ b/src/definitions.jl @@ -345,12 +345,11 @@ plan_irfft ############################################################################## """ - fftshift!(dest, src, [dim]) + fftshift!(dest, src, [dim]) Nonallocating version of [`fftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array. """ function fftshift!(dest, src, dim = 1:ndims(src)) - @assert size(dest)==size(src) s = ntuple(d -> d in dim ? div(size(dest,d),2) : 0, Val(ndims(dest))) circshift!(dest, src, s) end @@ -372,18 +371,14 @@ The output of `fftshift` is allocated. If one desires to store the output in a p """ fftshift -function fftshift(x, dim = 1:ndims(x)) - dest = similar(x) - fftshift!(dest, x, dim) -end +fftshift(x, dim = 1:ndims(x)) = fftshift!(similar(x), x, dim) """ - ifftshift!(dest, src, [dim]) + ifftshift!(dest, src, [dim]) Nonallocating version of [`ifftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array. """ function ifftshift!(dest, src, dim = 1:ndims(src)) - @assert size(dest)==size(src) s = ntuple(d -> d in dim ? -div(size(src,d),2) : 0, Val(ndims(src))) circshift!(dest, src, s) end @@ -405,10 +400,7 @@ The output of `ifftshift` is allocated. If one desires to store the output in a """ ifftshift -function ifftshift(x, dim = 1:ndims(x)) - dest = similar(x) - ifftshift!(dest, x, dim) -end +ifftshift(x, dim = 1:ndims(x)) = ifftshift!(similar(x), x, dim) ############################################################################## diff --git a/test/runtests.jl b/test/runtests.jl index 869c745..95c7c5d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -103,26 +103,26 @@ end @test @inferred(AbstractFFTs.fftshift([1 2 3])) == [3 1 2] @test @inferred(AbstractFFTs.fftshift([1, 2, 3])) == [3, 1, 2] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6])) == [6 4 5; 3 1 2] - a = [0 0 0] - b = [0, 0, 0] - c = [0 0 0; 0 0 0] - @test (AbstractFFTs.fftshift!(a, [1 2 3]); a == [3 1 2]) - @test (AbstractFFTs.fftshift!(b, [1, 2, 3]); b == [3, 1, 2]) + a = [0 0 0] + b = [0, 0, 0] + c = [0 0 0; 0 0 0] + @test (AbstractFFTs.fftshift!(a, [1 2 3]); a == [3 1 2]) + @test (AbstractFFTs.fftshift!(b, [1, 2, 3]); b == [3, 1, 2]) @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6]); c == [6 4 5; 3 1 2]) @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], 1)) == [4 5 6; 1 2 3] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], ())) == [1 2 3; 4 5 6] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], (1,2))) == [6 4 5; 3 1 2] @test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], 1:2)) == [6 4 5; 3 1 2] - @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3]) - @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6]) - @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [6 4 5; 3 1 2]) - @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1:2); c == [6 4 5; 3 1 2]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [6 4 5; 3 1 2]) + @test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1:2); c == [6 4 5; 3 1 2]) @test @inferred(AbstractFFTs.ifftshift([1 2 3])) == [2 3 1] @test @inferred(AbstractFFTs.ifftshift([1, 2, 3])) == [2, 3, 1] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6])) == [5 6 4; 2 3 1] - @test (AbstractFFTs.ifftshift!(a, [1 2 3]); a == [2 3 1]) + @test (AbstractFFTs.ifftshift!(a, [1 2 3]); a == [2 3 1]) @test (AbstractFFTs.ifftshift!(b, [1, 2, 3]); b == [2, 3, 1]) @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6]); c == [5 6 4; 2 3 1]) @@ -130,10 +130,10 @@ end @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], ())) == [1 2 3; 4 5 6] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], (1,2))) == [5 6 4; 2 3 1] @test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], 1:2)) == [5 6 4; 2 3 1] - @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3]) - @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6]) - @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [5 6 4; 2 3 1]) - @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1:2); c == [5 6 4; 2 3 1]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [5 6 4; 2 3 1]) + @test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1:2); c == [5 6 4; 2 3 1]) end @testset "FFT Frequencies" begin