Efficiently subset a 2D numpy array iteratively

Question

I have a large image on which I want to perform an operation from a moving window over the whole window. Here is a reproducible example:

Given an array named image with shape (5, 5), I need to extract subsets from the array in 3x3 windows.

import numpy as np

# example dada
image = np.array([[1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5], [1,2,3,4,5]])

Out[1]:

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

For a 3x3 window the first subset is:

# first iteration
window_size = 3
image[0:window_size, 0:window_size] # from 1st to 3th row and from 1st to 3th col

Out[2]:

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

So I can access to the different subsets using a nested loop:

for i in range(0, image.shape[0]-window_size+1):
   for j in range(0, image.shape[1]-window_size+1):
      a = (j,i)     # top left value
      b = (j,i+window_size)   # top right value
      c = (j+window_size,i)   # bottom left value
      d = (j+window_size,i+window_size) # bottom right value
      print('Window position', a,b,c,d)
      subset = image[i:i+window_size, j:j+window_size]
      print(subset)

Is there a more efficient way to perform this operation and avoid doing these two loops?

Answer 1

Thanks to @CrisLuengo and @mozway comments I can answer my own question. Indeed the dedicated function np.lib.stride_tricks.sliding_window_view is more efficient than a nested loop. Here is a small timing test with a window size of 3x3 for arias of 100x100, 1000x1000 and 10000x10000.

import numpy as np
import time

image = np.random.random((10000, 10000))
window_size = 3

# nested loop
start = time.time()
subsets = []
for i in range(0, image.shape[0]-window_size+1):
   for j in range(0, image.shape[1]-window_size+1):
      subset = image[i:i+window_size, j:j+window_size]
      subsets.append(subset)
end = time.time()
print(end - start)

# dedicated function
start = time.time()
subset = np.lib.stride_tricks.sliding_window_view(image, (window_size, window_size))
subsets = subset.reshape(subset.shape[0]*subset.shape[1], window_size, window_size)
end = time.time()
print(end - start)

However, there may be an even more efficient solution using scipy.ndimage.generic_filter. This function allows to apply a user-defined function to the different positions of the window. If I find a better solution with this approach I will update my answer.