Including the iterator in numpy sum function

user_13702483 0
Viewed 34

So I have the following function

def get_a(x):
    start = timer()
    A = np.zeros(len(x))
    for k in range(len(x)):
        Ak = 0
        Ak = np.sum(x*np.exp(-2j*pi*(k*1/len(x))))
        # for i in range(len(x)):
            # Ak+= x[i]*np.exp(-2j*pi*(k*i/len(x)))
        Ak = norm(Ak)
        A[k] = Ak
        
    end = timer()
    print(end - start) 
    return A

The commented part of the function is what I want to actually do, but it takes way too much time iterating over the array instead of using numpy's sum function, but the problem is that I need to somehow access the "i" value of the sum. In the line

Ak = np.sum(x*np.exp(-2j*pi*(k*1/len(x))))

What I need to replace is the "1" for the sum index. Is there any way to do it?

numpy sum python
1 Answers

Assuming the arrays mentioned are rank 1, you can produce a numpy array of the appropriate length with numpy.arange. Consider

Ak = np.sum(x*np.exp(-2j*pi*(k*np.arange(len(x))/len(x))))

There's some more refactoring that can be done. In particular, the whole fraction in the exponent can be replaced with an appropriate numpy.linspace call to get the appropriate ranges, but that should get you started vectorizing your equation.

user_22886590
Related