Question

is there a way to simplify a 3-nested loop over a 3d-field? The code look like:

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]

the shape of A is (N,N,N) and freq is a special iteration with the length N.
Maybe there is a numpy-tool to perform this code, cause this needs to long.

Answer 1

You can reduce time consumption by 20% with this simple optimization:

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]

Use dis module to check byte code and timeit to measure speed of your algorithm. Below examples on how your code has changed.

Every time you write int(k+0.5) you get following byte code:

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

It’s better to calculate it once, so the next call can be much faster:

     75 LOAD_FAST                2 (k)

Same story for N/2.0. Instead of having this in every iteration:

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

you can just use pre-calculated NN:

     78 LOAD_GLOBAL              4 (NN)