Hi - I have a scenario where I’m working with a large number of 3d arrays and want to efficiently calculate FFT’s for each column. I’ve written a few test functions that iterate in various ways across columns in calculating the associate 1d FFT’s shown below – along with test code. The slicing approach works fastest though reshaping comes close (much better memory and allocs). In digging into FFTW, it appears to allow ability to create a plan to do essentially the same using fftw_plan_many_dft. Does FFTW.jl support ability to use?
using FFTW
using BenchmarkTools
function scdfft(data::Array{ComplexF64, 3})
#get size tuple
stple = size(data)
#create data container
S = Array{ComplexF64}(undef, stple)
#fft across the rows, i.e. each column
for idx2 in 1:stple[2]
for idx3 in 1:stple[3]
@inbounds S[:, idx2, idx3] = fftshift(fft(@view data[:,idx2,idx3]))
end
end
return S
end
function scdfftslice(data::Array{ComplexF64, 3})
#create data container
S = Array{ComplexF64}(undef, size(data))
#fft down the columns
for (idx, col) in enumerate(eachslice(data, dims=2))
@inbounds S[:, idx, :] = fftshift(fft(col, 1), 1)
end
return S
end
function scdfftreshape(data::Array{ComplexF64, 3})
#get size tuple
stple = size(data)
#fft down the columns after reshaping
temp= fftshift(fft(reshape(data, (stple[1],:)), 1), 1)
return reshape(temp, stple)
end
Testing performance…
tst = rand(ComplexF64, (1024, 128,128));
#compile and confirm
scdfft(tst) == scdfftslice(tst) == scdfftreshape(tst) ? println("good") : error("bad")
good
@benchmark scdfft($tst)
Time (mean ± σ): 345.378 ms ± 14.753 ms ┊ GC (mean ± σ): 4.57% ± 0.13%
Memory estimate: 776.75 MiB, allocs estimate: 98
@benchmark scdfftslice($tst)
Time (mean ± σ): 207.515 ms ± 6.647 ms ┊ GC (mean ± σ): 5.50% ± 0.14%
Memory estimate: 768.04 MiB, allocs estimate: 898.
@benchmark scdfftreshape($tst)
Time (mean ± σ): 269.400 ms ± 11.756 ms ┊ GC (mean ± σ): 8.55% ± 2.11%
Memory estimate: 512.00 MiB, allocs estimate: 11.