Skip to content

Commit a7ebdac

Browse files
author
Grigorii Starkov
committed
Adds nonallocating versions of fftshift and ifftshift with the
corresponding tests. Fixes #62.
1 parent 4f630d6 commit a7ebdac

File tree

2 files changed

+39
-0
lines changed

2 files changed

+39
-0
lines changed

src/definitions.jl

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -364,6 +364,17 @@ function fftshift(x, dim = 1:ndims(x))
364364
circshift(x, s)
365365
end
366366

367+
"""
368+
fftshift!(dest, src, [dim])
369+
370+
Nonallocating version of [`fftshift`](@ref). Stores the result of the shift of the `src` array into the `dest` array.
371+
"""
372+
function fftshift!(dest, src, dim = 1:ndims(src))
373+
@assert size(dest)==size(src)
374+
s = ntuple(d -> d in dim ? div(size(dest,d),2) : 0, Val(ndims(dest)))
375+
circshift!(dest, src, s)
376+
end
377+
367378
"""
368379
ifftshift(x, [dim])
369380
@@ -384,6 +395,17 @@ function ifftshift(x, dim = 1:ndims(x))
384395
circshift(x, s)
385396
end
386397

398+
"""
399+
ifftshift!(dest, src, [dim])
400+
401+
Nonallocating version of [`ifftshift`](@ref). Stores the result of the shift of the `src` array into the `dest array`.
402+
"""
403+
function ifftshift!(dest, src, dim = 1:ndims(src))
404+
@assert size(dest)==size(src)
405+
s = ntuple(d -> d in dim ? -div(size(src,d),2) : 0, Val(ndims(src)))
406+
circshift!(dest, src, s)
407+
end
408+
387409
##############################################################################
388410

389411

test/runtests.jl

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,20 +103,37 @@ end
103103
@test @inferred(AbstractFFTs.fftshift([1 2 3])) == [3 1 2]
104104
@test @inferred(AbstractFFTs.fftshift([1, 2, 3])) == [3, 1, 2]
105105
@test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6])) == [6 4 5; 3 1 2]
106+
a = [0 0 0]
107+
b = [0, 0, 0]
108+
c = [0 0 0; 0 0 0]
109+
@test (AbstractFFTs.fftshift!(a, [1 2 3]); a == [3 1 2])
110+
@test (AbstractFFTs.fftshift!(b, [1, 2, 3]); b == [3, 1, 2])
111+
@test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6]); c == [6 4 5; 3 1 2])
106112

107113
@test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], 1)) == [4 5 6; 1 2 3]
108114
@test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], ())) == [1 2 3; 4 5 6]
109115
@test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], (1,2))) == [6 4 5; 3 1 2]
110116
@test @inferred(AbstractFFTs.fftshift([1 2 3; 4 5 6], 1:2)) == [6 4 5; 3 1 2]
117+
@test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3])
118+
@test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6])
119+
@test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [6 4 5; 3 1 2])
120+
@test (AbstractFFTs.fftshift!(c, [1 2 3; 4 5 6], 1:2); c == [6 4 5; 3 1 2])
111121

112122
@test @inferred(AbstractFFTs.ifftshift([1 2 3])) == [2 3 1]
113123
@test @inferred(AbstractFFTs.ifftshift([1, 2, 3])) == [2, 3, 1]
114124
@test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6])) == [5 6 4; 2 3 1]
125+
@test (AbstractFFTs.ifftshift!(a, [1 2 3]); a == [2 3 1])
126+
@test (AbstractFFTs.ifftshift!(b, [1, 2, 3]); b == [2, 3, 1])
127+
@test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6]); c == [5 6 4; 2 3 1])
115128

116129
@test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], 1)) == [4 5 6; 1 2 3]
117130
@test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], ())) == [1 2 3; 4 5 6]
118131
@test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], (1,2))) == [5 6 4; 2 3 1]
119132
@test @inferred(AbstractFFTs.ifftshift([1 2 3; 4 5 6], 1:2)) == [5 6 4; 2 3 1]
133+
@test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1); c == [4 5 6; 1 2 3])
134+
@test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], ()); c == [1 2 3; 4 5 6])
135+
@test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], (1,2)); c == [5 6 4; 2 3 1])
136+
@test (AbstractFFTs.ifftshift!(c, [1 2 3; 4 5 6], 1:2); c == [5 6 4; 2 3 1])
120137
end
121138

122139
@testset "FFT Frequencies" begin

0 commit comments

Comments
 (0)