/
fftshift.jl
83 lines (67 loc) · 2.09 KB
/
fftshift.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
"""
ft_center_diff(s [, dims])
Return the shifts required to center dimensions `dims` at the respective
Fourier centers.
This function is internally used by [`ShiftedArrays.fftshift`](@ref) and
[`ShiftedArrays.ifftshift`](@ref).
# Examples
```jldoctest
julia> ShiftedArrays.ft_center_diff((4, 5, 6), (1, 2)) # Fourier center is at (2, 3, 0)
(2, 2, 0)
julia> ShiftedArrays.ft_center_diff((4, 5, 6), (1, 2, 3)) # Fourier center is at (2, 3, 4)
(2, 2, 3)
```
"""
function ft_center_diff(s::NTuple{N, T}, dims=ntuple(identity, Val(N))) where {N, T}
return ntuple(i -> i ∈ dims ? s[i] ÷ 2 : 0, N)
end
"""
fftshift(x [, dims])
Lazy version of `AbstractFFTs.fftshift(x, dims)`. Return a `CircShiftedArray`
where each given dimension is shifted by `N÷2`, where `N` is the size of
that dimension.
# Examples
```jldoctest
julia> ShiftedArrays.fftshift([1 0 0 0])
1×4 CircShiftedArray{Int64, 2, Matrix{Int64}}:
0 0 1 0
julia> ShiftedArrays.fftshift([1 0 0; 0 0 0; 0 0 0])
3×3 CircShiftedArray{Int64, 2, Matrix{Int64}}:
0 0 0
0 1 0
0 0 0
julia> ShiftedArrays.fftshift([1 0 0; 0 0 0; 0 0 0], (1,))
3×3 CircShiftedArray{Int64, 2, Matrix{Int64}}:
0 0 0
1 0 0
0 0 0
```
"""
function fftshift(x::AbstractArray{T, N}, dims=ntuple(identity, Val(N))) where {T, N}
return ShiftedArrays.circshift(x, ft_center_diff(size(x), dims))
end
"""
ifftshift(x [, dims])
Lazy version of `AbstractFFTs.ifftshift(x, dims)`. Return a `CircShiftedArray`
where each given dimension is shifted by `-N÷2`, where `N` is the size of
that dimension.
# Examples
```jldoctest
julia> ShiftedArrays.ifftshift([0 0 1 0])
1×4 CircShiftedArray{Int64, 2, Matrix{Int64}}:
1 0 0 0
julia> ShiftedArrays.ifftshift([0 0 0; 0 1 0; 0 0 0])
3×3 CircShiftedArray{Int64, 2, Matrix{Int64}}:
1 0 0
0 0 0
0 0 0
julia> ShiftedArrays.ifftshift([0 1 0; 0 0 0; 0 0 0], (2,))
3×3 CircShiftedArray{Int64, 2, Matrix{Int64}}:
1 0 0
0 0 0
0 0 0
```
"""
function ifftshift(x::AbstractArray{T, N}, dims=ntuple(identity, Val(N))) where {T, N}
return ShiftedArrays.circshift(x, map(-, ft_center_diff(size(x), dims)))
end