How to use fft to 2d vector data

I have some data like this:

function f(x)
    [sin(2*pi*x[1])*cos(2*pi*x[2]), cos(2*pi*x[1])*sin(2*pi*x[2])]
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][1]
        c[i]=A[i][2]        
    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