Modifying all the frequency components of an FFT to 0 EXCEPT the 0th, +-3rd, +-6th, etc

New to Julia
, ,
1

I have an array and I want to modify its FFT frequencies by an exponential factor except those that are a multiple of 3 (or more generally, M). For example, we can do this naively as:

N = 256
M = 3
ψ = rand(N)
fft!(ψ)
for i = 2:N/2+1
    i = Int(i)
    if (i - 1) % M != 0
       ψ[i] *= exp(-*abs(ψ[i])) # + component
        ψ[N - i + 2] *= exp(-abs(ψ[N - i + 2])) # - component
    end
end 
ifft!(ψ)

I am aware of fftshift and the FFTView.jl package, but I can’t quite figure out how to do the indexing in a more elegant way (i.e. a 1 or 2 liner). This operation is done millions or billions of times in my production code so I’d like to do it more elegantly.

2

Don’t know anything about FFT unfortunately, but you might want to write your range as i = 2:N÷2 + 1 which saves you the Int(i) call on every iteration:

help?> ÷
"÷" can be typed by \div<tab>

search: ÷

  div(x, y)
  ÷(x, y)

  The quotient from Euclidean division. Computes x/y, truncated to an integer.
3

Instead of looping over a range with step 1, consider looping over a range with step M and remove the if statement.

Is the data ψ real? If so, consider using rfft and save half the computation and memory.

4

if N and M don’t change often, you can store the indices in an array (if they do and N is big, it might be a lot of allocations)

ind = [i for i in 2:(N÷2+1) if (i-1)%M != 0]
@views ψ[ind] .*= exp.(-abs.(ψ[ind]))

with @views to avoid allocating new arrays

5

Thank you, I ended up with:

ind_= [i for i in 2:(N÷2+1) if (i-1)%M != 0]
ind = sort([ind_p; N.-ind.+2])
@views ψ[ind] .*= exp.(-abs.(ψ[ind]))