Python: 嵌套循环遍历3D

Question

有没有办法简化三维场上的三嵌套循环？代码如下所示：

from itertools import product
for kx, ky, kz in product(freq, freq, freq):
    k = np.sqrt(kx**2+ky**2+kz**2)
    if int(k+0.5) < N/2.0:
        yaxes[field][int(k+0.5)] += A[kx][ky][kz]

A的形状是（N，N，N），freq是长度为N的特殊迭代。 也许有一个numpy工具来执行这段代码，因为这需要很长时间。你知道吗

Answer 1

通过以下简单的优化，您可以将时间消耗减少20%：

NN = N/2.0
for kx, ky, kz in product(freq, freq, freq):
    k = int(np.sqrt(kx**2 + ky**2 + kz**2) + 0.5)
    if k < NN:
        yaxes[field][k] += A[kx][ky][kz]

使用dis模块检查字节码并计时以测量算法的速度。下面的例子说明了你的代码是如何变化的。你知道吗

每次写入int（k+0.5）时，都会得到以下字节码：

     65 LOAD_GLOBAL              3 (int) 
     68 LOAD_FAST                2 (k) 
     71 LOAD_CONST               3 (0.5) 
     74 BINARY_ADD           
     75 CALL_FUNCTION            1 

最好只计算一次，这样下一次调用可以快得多：

     75 LOAD_FAST                2 (k)

N/2.0也是这样。而不是在每次迭代中都有这个：

     78 LOAD_GLOBAL              4 (N) 
     81 LOAD_CONST               4 (2.0) 
     84 BINARY_TRUE_DIVIDE   

只需使用预先计算的NN：

     78 LOAD_GLOBAL              4 (NN)