Vectorized moving window on 2D array in numpy

If you can express the function

f(in[i][j],in[i+1][j],in[i-1][j],in[i][j+1],in[i][j-1],…)

as a linear operator, you could use scipy's signal.convolve2d function to do exactly that. For instance, say you have an 50x50 array, A, and you want to calculate a second array B where each of its element b[ij] is the average over a[i,j], a[(i-1),j], a[i,(j-1)], a[(i-1),(j-1)] from the array A. You could do that simply doing :

A = # your first array
B = numpy.ones((2,2))/4
C = scipy.signal.convolve2d(A,B, 'valid')

When the convolution is performed, the array B "slides" across A, multiplying the corresponding elements and summing up the result. Because of border effects, you must be careful when using the resulting array C. Here, C is of shape 49x49, because of the 'valid' argument in convolve2d, to discard the first row and column since they contain border effects. If you wanted to have a 50x50 array, without discarding, you would swap that argument for 'same'

EDIT: Perhaps if you could tell me more about that function you need, I could help you more specifically in turning it into an array that would be used to do the 2D convolution.

Hope that helps!


You can use the rolling window technique as explained here, here and here, but for 2D array.

The source code for 2D rolling window in NumPy:

# Rolling window for 2D arrays in NumPy
import numpy as np

def rolling_window(a, shape):  # rolling window for 2D array
    s = (a.shape[0] - shape[0] + 1,) + (a.shape[1] - shape[1] + 1,) + shape
    strides = a.strides + a.strides
    return np.lib.stride_tricks.as_strided(a, shape=s, strides=strides)

a = np.array([[0,  1,  2,  3,  4,  5],
              [6,  7,  8,  9, 10,  11],
              [12, 13, 14, 15, 7,   8],
              [18, 19, 20, 21, 13, 14],
              [24, 25, 26, 27, 19, 20],
              [30, 31, 32, 33, 34, 35]], dtype=np.int)
b = np.arange(36, dtype=np.float).reshape(6,6)
present = np.array([[7,8],[13,14],[19,20]], dtype=np.int)
absent  = np.array([[7,8],[42,14],[19,20]], dtype=np.int)

found = np.all(np.all(rolling_window(a, present.shape) == present, axis=2), axis=2)
print(np.transpose(found.nonzero()))
found = np.all(np.all(rolling_window(b, present.shape) == present, axis=2), axis=2)
print(np.transpose(found.nonzero()))
found = np.all(np.all(rolling_window(a, absent.shape) == absent, axis=2), axis=2)
print(np.transpose(found.nonzero()))

Array present is occurred in array a two times on [1,1] and [2,4].

More examples in my CoLab notebook "Rolling window on NumPy arrays without for loops".