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

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.

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.

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.

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

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]))