Search code examples
pythonnumpyperformancenumerical

Speed up a program that calculate the average of the neighbors in a huge array


I have a problem with the speed of my program. I want to calculate the average of four neighbors in a huge array. Here is a part of my code. Do you have any ideas how to change the last line? Or should I use another array?

for a in np.arange(100000):
    for x in np.arange(size):
        for y in np.arange(size):
            if unchangeableflag[x*size+y] == 0:
                vnew[x*size+y] = (v[(x+1)*size+y] + v[(x-1)*size+y] + v[x*size+y+1] + v[x*size+y-1]) / 4.0

Solution

  • You would not need the loop at all. Assuming v, vnew and unchangeableflag are 1-d arrays with size*size entries, you can do

    v = v.reshape(size, size)
    vnew = vnew.reshape(size, size)
    unchangeableflag = unchangeableflag.reshape(size, size)
    average = v[1:-1, 2:]
    average += v[1:-1, :-2] 
    average += v[2:, 1:-1]
    average += v[-2:, 1:-1]
    average /= 4.0
    vnew[1:-1, 1:-1][unchangeableflag[1:-1, 1:-1] == 0] = average
    

    But what are you actually trying to achieve? This looks suspiciously like you could get away with some application of the discrete Laplacian.

    (Note that this assumes that v contains floating point numbers. If the dtype of `v´ is sime integral type, you need a slight modification.)