How to apply fft to an array of static arrays?

General Usage
,
1

I’d like to be able to do something like

using StaticArrays, FFTW
A = [@SVector rand(ComplexF64, 3) for m ∈ 1:512]
fft(A)

but FFTW.jl throws an error when I try it.

Due to performance and ergonomic reasons, I really want to use an array of static vectors, instead of an array with one extra dimension. But just to be clear, I’m trying to perform an operation analogous to

A = rand(ComplexF64, 3, 512)
fft(A, dims=2)

Is there a way to accomplish this with an array of SArrays? I’d like to also do this on the GPU.

Thanks in advance!

2

There is a @SMatrix macro for 2-dimensional static arrays.

julia> using StaticArrays, FFTW

julia> A = @SMatrix rand(ComplexF64, 3, 512)
3×512 SMatrix{3, 512, ComplexF64, 1536} with indices SOneTo(3)×SOneTo(512):
 0.700458+0.261335im   0.380832+0.435067im  …   0.94796+0.0524721im
 0.789127+0.51047im    0.790958+0.100901im     0.486345+0.982996im
 0.813353+0.462645im  0.0897808+0.122944im     0.550222+0.546134im

julia> fft(A, 2)
3×512 Matrix{ComplexF64}:
   261.6+254.863im  0.532757+2.99622im  …   -4.6475+5.76742im
 247.911+260.101im  -2.50283+2.4603im      -3.08027+6.46999im
 251.814+255.649im  -8.23491+4.1779im      -2.99529-13.4493im

julia> fft(A, 2)[:,1:1] ≈  sum(A, dims=2)
true

Someone else will have to answer for GPUs.

1 Like
3

Maybe you want fft(reinterpret(reshape, ComplexF64, A), 2)? That ought to be a little quicker than fft(stack(A), 2). You can reinterpret back afterwards (or call eachcol).

4 Likes
4

Thanks for the answer! Nonetheless, I really want to stick with arrays of static arrays. I’ve edited the post to make this point clearer

5

Thanks, I guess this is it! I didn’t know about this three argument reinterpret. Is this expected to work on GPU as well? I currently do not have access to one to try things out.

1 Like