Suppose I have the following loop:

N=5
a=np.zeros((N,N))

for i in range(N):
  for j in range(N):
      for k in range(N):
        for l in range(N):
          a[i,j] =np.exp(1j*(2*np.pi/N*i*k 2*np.pi*j*l))

How can I optimize this? I'm out of ideas

CodePudding user response：

import numpy as np

x = np.arange(N)
i = x[:, None, None, None]
j = x[None, :, None, None]
k = x[None, None, :, None]
l = x[None, None, None, :]

out = np.exp(1j*(2*np.pi/N*i*k   2*np.pi*j*l)).sum(axis=(2, 3))

# >>> np.allclose(a, out)
# True