Numpy: How to slice or split 2D subsections of 2D array

Question

I have a 2D array. For example:

ary = np.arange(24).reshape(6,4)

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]
 [16 17 18 19]
 [20 21 22 23]]

I want to break this into smaller 2D arrays, each 2x2, and compute the square root of the sum of each. I actually want to use arbitrary sized sub-arrays, and compute arbitrary functions of them, but I think this question is easier to ask with concrete operations and concrete array sizes, so in this example starting with a 6x4 array and computing the square root of sums of 2x2 sub-arrays, the final result would be a 3x2 array, as follows:

[[3.16, 4.24]    # math.sqrt(0+1+4+5)    , math.sqrt(2+3+6+7)
 [6.48, 7.07]    # math.sqrt(8+9+12+13)  , math.sqrt(10+11+14+15)
 [8.60, 9.05]]   # math.sqrt(16+17+20+21), math.sqrt(18+19+22+23)

How can I slice, or split, or do some operation to perform some computation on 2D sub-arrays?

Here is a working, inefficient example of what I'm trying to do:

import numpy as np

a_height = 6
a_width = 4
a_area = a_height * a_width
a = np.arange(a_area).reshape(a_height, a_width)

window_height = 2
window_width = 2

b_height = a_height // window_height
b_width = a_width // window_width
b_area = b_height * b_width
b = np.zeros(b_area).reshape(b_height, b_width)

for i in range(b_height):
    for j in range(b_width):
        b[i, j] = a[i * window_height:(i + 1) * window_height, j * window_width:(j + 1) * window_width].sum()
b = np.sqrt(b)

print(b)

# [[3.16227766 4.24264069]
#  [6.4807407  7.07106781]
#  [8.60232527 9.05538514]]

Answer 1

In [2]: ary = np.arange(24).reshape(6,4)
In [3]: ary
Out[3]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15],
       [16, 17, 18, 19],
       [20, 21, 22, 23]])

While I recommended moving-windows based on as_strided, we can also divide the array into 'blocks' with reshape and transpose:

In [4]: ary.reshape(3,2,2,2).transpose(0,2,1,3)
Out[4]: 
array([[[[ 0,  1],
         [ 4,  5]],

        [[ 2,  3],
         [ 6,  7]]],


       [[[ 8,  9],
         [12, 13]],

        [[10, 11],
         [14, 15]]],


       [[[16, 17],
         [20, 21]],

        [[18, 19],
         [22, 23]]]])
In [5]: np.sqrt(_.sum(axis=(2,3)))
Out[5]: 
array([[3.16227766, 4.24264069],
       [6.4807407 , 7.07106781],
       [8.60232527, 9.05538514]])

While the transpose makes it easier to visual the blocks that need to be summed, it isn't necessary:

In [7]: np.sqrt(ary.reshape(3,2,2,2).sum(axis=(1,3)))
Out[7]: 
array([[3.16227766, 4.24264069],
       [6.4807407 , 7.07106781],
       [8.60232527, 9.05538514]])

np.lib.stride_tricks.sliding_window doesn't give us as much direct control as I thought, but

np.lib.stride_tricks.sliding_window_view(ary,(2,2))[::2,::2]

gives the same result as Out[4].

In [13]: np.sqrt(np.lib.stride_tricks.sliding_window_view(ary,(2,2))[::2,::2].sum(axis=(2,3)))
Out[13]: 
array([[3.16227766, 4.24264069],
       [6.4807407 , 7.07106781],
       [8.60232527, 9.05538514]])

[7] is faster.

In general, it can be done like this:

a_height = 15
a_width = 16
a_area = a_height * a_width
a = np.arange(a_are).reshape(a_height, a_width)

window_height = 3  # must evenly divide a_height
window_width = 4  # must evenly divide a_width

b_height = a_height // window_height
b_width = a_width // window_width

b = a.reshape(b_height, window_height, b_width, window_width).transpose(0,2,1,3)

# or, assuming you want sum or another function that takes `axis` argument
b = a.reshape(b_height, window_height, b_width, window_width).sum(axis=(1,3))