FOR loop performance improvement using numpy

Question

The following snippet of code takes 0.7s. I would like to improve the speed. Basically, the code adds all values in b with the same index as a, and stores them in the same position, found in a's index, but in a different array. So basically array a holds values ranging from 0-255, which represent the indexes of temp array.

a = np.random.randint(256, size=(40000,2))
b= np.arange(1280000).reshape(40000, 32)
temp = np.zeros((1,32,256,256)) 

for indx, pnt in enumerate(a):
    temp[0,:,pnt[0],pnt[1]] += b[indx,:]

Thanks.

Answer 1

This answer might be a bit left field, but I think it's not worth bending over backwards to try to vectorize.

Throw some numba on this problem instead;

from numba import jit

@jit
def compute_stuff(a, b):
    temp = np.zeros((1,32,256,256)) 
    for indx, pnt in enumerate(a):
        temp[0,:,pnt[0],pnt[1]] += b[indx,:]
    return temp

and it's at least a bit faster.

It also seems like you know that your a must be less than 256, so you can save memory and likely gain some performance by specifying the datatype for your array;

a = np.random.randint(256, size=(40000,2), dtype=np.uint8)