I have a large 2-dimensional numpy array, with every value being a 0 or a 1. I'd like to create a function which takes this array as an input, and returns a new array of the same size in which each element is based on the elements above, below and to either side. The returned array should have 0's stay 0, and each 1 will get +1 if there is a 1 to the north, +2 for a 1 to the right, +4 for a 1 below and +8 for a 1 to the left. These all stack, so a 1 surrounded by 1's should end up as a 17. Diagonals do not matter. This might also be faster with explicit bitwise operations (with 4 bits, each bit corresponding to a direction and whether there is a 1 or 0 in it).
I would like this operation to be done as quickly as possible. I played around with a for loop but it is too slow, and I don't understand masking in numpy well enough to use that.
I hope this can help. I start copying the original matrix, then add the contribution from each direction. For example if I have to add the contribute of the elements on the right, they may modify all the columns but the last one, thus I can write result[:,:-1] += m[:,1:]
. The last multiplication for m
ensures that the starting value of each cell to modify was one and not zero, as you required.
import numpy as np
def f(m):
result = np.copy(m)
# propagation from the four directions
result[1:,:] += m[:-1,:] # north
result[:,:-1] += 2 * m[:,1:] # est
result[:-1,:] += 4 * m[1:,:] # sud
result[:,1:] += 8 * m[:,:-1] # west
return result * m