使用numpy优化嵌套循环

Question

假设我有以下循环:

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

我如何优化这个?我没主意了

Answer 1

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