Vectorized 2-D moving window in numpy including edges

Question

I realize my question is fairly similar to Vectorized moving window on 2D array in numpy , but the answers there don't quite satisfy my needs.

Is it possible to do a vectorized 2D moving window (rolling window) which includes so-called edge effects? What would be the most efficient way to do this?

That is, I would like to slide the center of a moving window across my grid, such that the center can move over each cell in the grid. When moving along the margins of the grid, this operation would return only the portion of the window that overlaps the grid. Where the window is entirely within the grid, the full window is returned. For example, if I have the grid:

array([[1,2,3,4],
       [2,3,4,5],
       [3,4,5,6],
       [4,5,6,7]])

…and I want to sample each point in this grid using a 3x3 window centered at that point, the operation should return a series of arrays, or, ideally, a series of views into the original array, as follows:

array([[1,2],    array([[1,2,3],    array([[2,3,4],    array([[3,4],
       [2,3]])          [2,3,4]])          [3,4,5]])          [4,5]])

array([[1,2],    array([[1,2,3],    array([[2,3,4],    array([[3,4],
       [2,3],           [2,3,4],           [3,4,5],           [4,5],
       [3,4]])          [3,4,5]])          [4,5,6]])          [5,6]])

array([[2,3],    array([[2,3,4],    array([[3,4,5],    array([[4,5],
       [3,4],           [3,4,5],           [4,5,6],           [5,6],
       [4,5]])          [4,5,6]])          [5,6,7]])          [6,7]])

array([[3,4],    array([[3,4,5],    array([[4,5,6],    array([[5,6],
       [4,5]])          [4,5,6]])          [5,6,7]])          [6,7]])

Because I need to perform this operation many times, speed is important & the ideal solution would be a vectorized operation.

Answer 1

You could define a function that yields a generator and use that. The window would be the floor of the shape you want divided by 2 and the trick would be just indexing the array along that window as you move along the rows and columns.

def window(arr, shape=(3, 3)):
    # Find row and column window sizes
    r_win = np.floor(shape[0] / 2).astype(int)
    c_win = np.floor(shape[1] / 2).astype(int)
    x, y = arr.shape
     for i in range(x):
         xmin = max(0, i - r_win)
         xmax = min(x, i + r_win + 1)
         for j in range(y):
             ymin = max(0, j - c_win)
             ymax = min(y, j + c_win + 1)
             yield arr[xmin:xmax, ymin:ymax]

You could use this function like so:

arr = np.array([[1,2,3,4],
               [2,3,4,5],
               [3,4,5,6],
               [4,5,6,7]])
gen = window(arr)
next(gen)
array([[1, 2],
       [2, 3]])

Going through the generator produces all of the windows in your example.

It's not vectorized, but I'm not sure there is an existing vectorized function that returns different sized arrays. As @PaulPanzer points out you could pad your array to the size you need and use a np.lib.stride_tricks.as_strided to generate a view of the slices. Something like so:

def rolling_window(a, shape):
    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)

def window2(arr, shape=(3, 3)):
    r_extra = np.floor(shape[0] / 2).astype(int)
    c_extra = np.floor(shape[1] / 2).astype(int)
    out = np.empty((arr.shape[0] + 2 * r_extra, arr.shape[1] + 2 * c_extra))
    out[:] = np.nan
    out[r_extra:-r_extra, c_extra:-c_extra] = arr
    view = rolling_window(out, shape)
    return view

window2(arr, (3,3))
array([[[[ nan,  nan,  nan],
         [ nan,   1.,   2.],
         [ nan,   2.,   3.]],

        [[ nan,  nan,  nan],
         [  1.,   2.,   3.],
         [  2.,   3.,   4.]],

        [[ nan,  nan,  nan],
         [  2.,   3.,   4.],
         [  3.,   4.,   5.]],

        [[ nan,  nan,  nan],
         [  3.,   4.,  nan],
         [  4.,   5.,  nan]]],


       [[[ nan,   1.,   2.],
         [ nan,   2.,   3.],
         [ nan,   3.,   4.]],

        [[  1.,   2.,   3.],
         [  2.,   3.,   4.],
         [  3.,   4.,   5.]],

        [[  2.,   3.,   4.],
         [  3.,   4.,   5.],
         [  4.,   5.,   6.]],

        [[  3.,   4.,  nan],
         [  4.,   5.,  nan],
         [  5.,   6.,  nan]]],


       [[[ nan,   2.,   3.],
         [ nan,   3.,   4.],
         [ nan,   4.,   5.]],

        [[  2.,   3.,   4.],
         [  3.,   4.,   5.],
         [  4.,   5.,   6.]],

        [[  3.,   4.,   5.],
         [  4.,   5.,   6.],
         [  5.,   6.,   7.]],

        [[  4.,   5.,  nan],
         [  5.,   6.,  nan],
         [  6.,   7.,  nan]]],


       [[[ nan,   3.,   4.],
         [ nan,   4.,   5.],
         [ nan,  nan,  nan]],

        [[  3.,   4.,   5.],
         [  4.,   5.,   6.],
         [ nan,  nan,  nan]],

        [[  4.,   5.,   6.],
         [  5.,   6.,   7.],
         [ nan,  nan,  nan]],

        [[  5.,   6.,  nan],
         [  6.,   7.,  nan],
         [ nan,  nan,  nan]]]])

This version pads the edges with np.nan to avoid confusion with any other values in your array. It is about 3x faster with the given array than the window function, but I am not sure how having padded output will impact anything you want to do downstream.

Answer 2

This is not strictly an answer to your question since its not vectorized, but hopefully its a helpful benchmark for any other potential solutions (there's bound to be something in image processing libraries?)

Regardless, I've implemented the window as a loop which takes the average of the window with output into a new array. Inputs are an array and the window size +/- the current index. One version using straight Python and Numpy, the other compiled with numba.

def mw_mean(in_arr,out_arr,x_win,y_win):
    xn,yn = in_arr.shape
    for x in range(xn):
        xmin = max([0,x - x_win])
        xmax = min([xn, x + x_win + 1])
        for y in range(yn):
            ymin = max([0,y - y_win])
            ymax = min([yn, y + y_win + 1])

            out_arr[x,y] = in_arr[xmin:xmax, ymin:ymax].mean()
    return out_arr



@jit("i4[:,:](i4[:,:],i4[:,:],i4,i4)", nopython = True)
def mw_mean_numba(in_arr,out_arr,x_win,y_win):
    xn,yn = in_arr.shape
    for x in range(xn):
        xmin = max(0,x - x_win)
        xmax = min(xn, x + x_win + 1)
        for y in range(yn):
            ymin = max(0,y - y_win)
            ymax = min(yn, y + y_win + 1)

            out_arr[x,y] = in_arr[xmin:xmax, ymin:ymax].mean()
    return out_arr

This is tested against three different array sizes--your original test case and two larger ones (100x100 and 1000x1000):

a = np.array([[1,2,3,4], [2,3,4,5], [3,4,5,6], [4,5,6,7]])
b = np.random.randint(1,7, size = (100,100))
c = np.random.randint(1,7, size = (1000,1000))

aout,bout,cout = np.zeros_like(a),np.zeros_like(b),np.zeros_like(c)

x_win = 1
y_win = 1

Runtimes without compilation:

%timeit mw_mean(a,aout,x_win,y_win)
1000 loops, best of 3: 225 µs per loop

%timeit mw_mean(b,bout,x_win,y_win)
10 loops, best of 3: 137 ms per loop

%timeit mw_mean(c,cout,x_win,y_win)
1 loop, best of 3: 14.1 s per loop

Runtimes with compilation:

%timeit mw_mean_numba(a,aout,x_win,y_win)
1000000 loops, best of 3: 1.22 µs per loop

%timeit mw_mean_numba(b,bout,x_win,y_win)
1000 loops, best of 3: 550 µs per loop

%timeit mw_mean_numba(c,cout,x_win,y_win)
10 loops, best of 3: 55.1 ms per loop

Edit: a previous version of this was modifying the array in place, which is obviously a big no-no for a rolling window. Benchmarks remain unchanged.