# How to use fft to 2d vector data

I have some data like this:

function f(x)
[sin(2*pi*x)*cos(2*pi*x), cos(2*pi*x)*sin(2*pi*x)]
end
N=200
data = [f([i/N,j/N]) for i in 0:N-1, j in 0:N-1]


which is a matrix of size N\times N, with vector elements. Now I want to take fft in each slice of data so that the result has the same structure as data. I can do this by using the follow codes

using FFTW
function v_fft(A::Matrix{Vector{Float64}})
n=size(A,2)
b=Matrix{Float64}(undef, n, n)
c=Matrix{Float64}(undef, n, n)
@inbounds @simd for i in 1:n^2
b[i]=A[i]
c[i]=A[i]
end
fb=fft(b)
fc=fft(c)
result=Matrix{Vector{ComplexF64}}(undef, n, n)
@inbounds @simd for i in 1:n^2
result[i]=[fb[i],fc[i]]
end
result
end


v_fft(data) gives what I want. Is there any simple way to do that ?

In the nested format there is probably no easier way. Using instead an N x N x 2 matrix layout, you could use eachslice, broadcasting and stack:

julia> using Compat

julia> datamat = permutedims(stack(data), (2, 3, 1));

julia> size(datamat)
(200, 200, 2)

julia> res = stack(fft.(eachslice(datamat; dims = 3)));

julia> res[:, :, 1] == first.(v_fft(data))
true

julia> size(res)
(200, 200, 2)

2 Likes

Thank you! This works for me.

1 Like