Sliding Window over stacked np.array (Python)

Question

I am trying to create a multi-dimensional window across a stack np array and find the max of each of the windows. Below res is a stacked array with shape (3, 4, 4). I would like to have a window with 2x2 size through the axis. For example the first window will be (3,2,2) with values:

ideal_result = np.array([[13, 45], [1, 2], [11, 22], [1, 2], [1, 2], [1, 7]])

then max the window would be: np.max(ideal_result) = 45

This would be be the entire windows and construct a 2x2 with np.array([[45 67],[23 44]])

Set up:

a = np.array([[13, 45, 67, 4], [1, 2, 3, 4], [2, 3, 4, 6], [1, 23, 44, 1]])
b = np.array([[11, 22, 33, 57], [1, 2, 3, 4], [2, 3, 94, 6], [1, 23, 44, 1]])
c = np.array([[1, 2, 3, 4], [1, 7, 8, 9], [2, 3, 4, 76], [1, 23, 44, 1]])
res = np.stack((a, b, c))
print(np.shape(res))

Code Attempted:

import numpy as np
v = np.lib.stride_tricks.as_strided(res, shape=(3, 2, 2), strides=(3, 2, 2))

Answer 1

It usually helps to think about how the flattened array looks to get the strides right:

res.flatten()

output:

array([13, 45, 67,  4,  1,  2,  3,  4,  2,  3,  4,  6,  1, 23, 44,  1, 11,
       22, 33, 57,  1,  2,  3,  4,  2,  3, 94,  6,  1, 23, 44,  1,  1,  2,
        3,  4,  1,  7,  8,  9,  2,  3,  4, 76,  1, 23, 44,  1])

the dtype is int64 which is 8 bytes, so the stride between each consecutive element is 8. I'll mark which elements should be in the first window with the stride:

        0   1           4   5                                          16
array([13, 45, 67,  4,  1,  2,  3,  4,  2,  3,  4,  6,  1, 23, 44,  1, 11,
        
       17          20  21                                          32  33
       22, 33, 57,  1,  2,  3,  4,  2,  3, 94,  6,  1, 23, 44,  1,  1,  2,

               36  37
        3,  4,  1,  7,  8,  9,  2,  3,  4, 76,  1, 23, 44,  1])

Can you see the pattern?

We have the following values and strides for each dimension:

x| values: 13->45, 11->22, 1->2, ... stride: 1
y| values: 13->1, 45->2, 11->2, ... stride: 4
z| values: 13->11, 45->22, 11->1, ... stride: 16

np.lib.stride_tricks.as_strided(res, shape=(3, 2, 2), strides=(8 * 16, 8 * 4, 8 * 1))

output:

array([[[13, 45],
        [ 1,  2]],

       [[11, 22],
        [ 1,  2]],

       [[ 1,  2],
        [ 1,  7]]])

This is just one window and we want 4 of them, the jump between each window is 8*2 bytes in x direction and 8*8 bytes in y direction.

windows = np.lib.stride_tricks.as_strided(res, shape=(2, 2, 3, 2, 2), strides=(8 * 8, 8 * 2, 8 * 16, 8 * 4, 8 * 1))
windows.max(axis=(2, 3, 4))

output:

array([[45, 67],
       [23, 94]])

Assuming that the stride is the same as the kernel size (as in conventional 2D max pooling).