2d numpy array, making each value the sum of the 3x3 square it is centered at

Question

I have a square 2D numpy array, A, and an array of zeros, B, with the same shape.

For every index (i, j) in A, other than the first and last rows and columns, I want to assign to B[i, j] the value of np.sum(A[i - 1:i + 2, j - 1:j + 2].

Example:

A =
array([[0, 0, 0, 0, 0],
       [0, 1, 0, 1, 0],
       [0, 1, 1, 0, 0],
       [0, 1, 0, 1, 0],
       [0, 0, 0, 0, 0])

B =
array([[0, 0, 0, 0, 0],
       [0, 3, 4, 2, 0],
       [0, 4, 6, 3, 0],
       [0, 3, 4, 2, 0],
       [0, 0, 0, 0, 0])

Is there an efficient way to do this? Or should I simply use a for loop?

Answer 1

There is a clever (read "borderline smartass") way to do this with np.lib.stride_tricks.as_strided. as_strided allows you to create views into your buffer that simulate windows by adding another dimension to the view. For example, if you had a 1D array like

>>> x = np.arange(10)
>>> np.lib.stride_tricks.as_strided(x, shape=(3, x.shape[0] - 2), strides=x.strides * 2)
array([[0, 1, 2, 3, 4, 5, 6, 7],
       [1, 2, 3, 4, 5, 6, 7, 8],
       [2, 3, 4, 5, 6, 7, 8, 9]])

Hopefully it is clear that you can just sum along axis=0 to get the sum of each size 3 window. There is no reason you couldn't extrend that to two or more dimensions. I've written the shape and index of the previous example in a way that suggests a solution:

A = np.array([[0, 0, 0, 0, 0],
              [0, 1, 0, 1, 0],
              [0, 1, 1, 0, 0],
              [0, 1, 0, 1, 0],
              [0, 0, 0, 0, 0]])
view = np.lib.stride_tricks.as_strided(A,
    shape=(3, 3, A.shape[0] - 2, A.shape[1] - 2),
    strides=A.strides * 2
)
B[1:-1, 1:-1] = view.sum(axis=(0, 1))

Summing along multiple axes simultaneously has been supported in np.sum since v1.7.0. For older versions of numpy, just sum repeatedly (twice) along axis=0.

Filling in the edges of B is left as an exercise for the reader (since it's not really part of the question).

As an aside, the solution here is a one-liner if you want it to be. Personally, I think anything with as_strided is already illegible enough, and doesn't need any further obfuscation. I'm not sure if a for loop is going to be bad enough performance-wise to justify this method in fact.

For future reference, here is a generic window-making function that can be used to solve this sort of problem:

def window_view(a, window=3):
    """
    Create a (read-only) view into `a` that defines window dimensions.

    The first ``a.ndim`` dimensions of the returned view will be sized according to `window`.
    The remaining ``a.ndim`` dimensions will be the original dimensions of `a`, truncated by `window - 1`.

    The result can be post-precessed by reducing the leading dimensions. For example, a multi-dimensional moving average could look something like ::


         window_view(a, window).sum(axis=tuple(range(a.ndim))) / window**a.ndim

    If the window size were different for each dimension (`window` were a sequence rather than a scalar), the normalization would be ``np.prod(window)`` instead of ``window**a.ndim``.

    Parameters
    -----------
    a : array-like
        The array to window into. Due to numpy dimension constraints, can not have > 16 dims.
    window :
        Either a scalar indicating the window size for all dimensions, or a sequence of length `a.ndim` providing one size for each dimension.

    Return
    ------
    view : numpy.ndarray
         A read-only view into `a` whose leading dimensions represent the requested windows into `a`.
         ``view.ndim == 2 * a.ndim``.
    """
    a = np.array(a, copy=False, subok=True)
    window = np.array(window, copy=False, subok=False, dtype=np.int)
    if window.size == 1:
        window = np.full(a.ndim, window)
    elif window.size == a.ndim:
        window = window.ravel()
    else:
        raise ValueError('Number of window sizes must match number of array dimensions')

    shape = np.concatenate((window, a.shape))
    shape[a.ndim:] -= window - 1
    strides = a.strides * 2

    return np.lib.stride_tricks.as_strided(a, shake=shape, strides=strides)

Answer 2

I have found no 'simple' ways of doing this. But here are two ways:

---

• Still involves a for loop

# Basically, get the sum for each location and then pad the result with 0's
B = [[np.sum(A[j-1:j+2,i-1:i+2]) for i in range(1,len(A)-1)] for j in range(1,len(A[0])-1)]
B = np.pad(B, ((1,1)), "constant", constant_values=(0))

---

• Is longer but no for loops (this will be a lot more efficient on big arrays):

# Roll basically slides the array in the desired direction
A_right = np.roll(A, -1, 1)
A_left = np.roll(A, 1, 1)
A_top = np.roll(A, 1, 0)
A_bottom = np.roll(A, -1, 0)
A_bot_right = np.roll(A_bottom, -1, 1)
A_bot_left = np.roll(A_bottom, 1, 1)
A_top_right = np.roll(A_top, -1, 1)
A_top_left = np.roll(A_top, 1, 1)

# After doing that, you can just add all those arrays and these operations
# are handled better directly by numpy compared to when you use for loops
B = A_right + A_left + A_top + A_bottom + A_top_left + A_top_right + A_bot_left + A_bot_right + A

# You can then return the edges to 0 or whatever you like
B[0:len(B),0] = 0
B[0:len(B),len(B[0])-1] = 0
B[0,0:len(B)] = 0
B[len(B[0])-1,0:len(B)] = 0