How to evaluate the sum of values within array blocks

Question

I have data array, with shape 100x100. I want to divide it into 5x5 blocks, and each block has 20x20 grids. The value of each block I want is the sum of all values in it.

Is there a more elegant way to accomplish it?

x = np.arange(100)
y = np.arange(100)
X, Y = np.meshgrid(x, y)
Z = np.cos(X)*np.sin(Y)
Z_new = np.zeros((5, 5))
for i in range(5):
  for j in range(5):
    Z_new[i, j] = np.sum(Z[i*20:20+i*20, j*20:20+j*20])

This is based on index, how if based on x?

x = np.linspace(0, 1, 100)
y = np.linspace(0, 1, 100)
X, Y = np.meshgrid(x, y)
Z = np.cos(X)*np.sin(Y)
x_new = np.linspace(0, 1, 15)
y_new = np.linspace(0, 1, 15)

Z_new?

Answer 1

Simply reshape splitting each of those two axes into two each with shape (5,20) to form a 4D array and then sum reduce along the axes having the lengths 20, like so -

Z_new = Z.reshape(5,20,5,20).sum(axis=(1,3))

Functionally the same, but potentially faster option with np.einsum -

Z_new = np.einsum('ijkl->ik',Z.reshape(5,20,5,20))

Generic block size

Extending to a generic case -

H,W = 5,5 # block-size
m,n = Z.shape
Z_new = Z.reshape(H,m//H,W,n//W).sum(axis=(1,3))

With einsum that becomes -

Z_new = np.einsum('ijkl->ik',Z.reshape(H,m//H,W,n//W))

To compute average/mean across blocks, use mean instead of sum method.

Generic block size and reduction operation

Extending to use reduction operations that have ufuncs supporting multiple axes parameter with axis for reductions, it would be -

def blockwise_reduction(a, height, width, reduction_func=np.sum):
    m,n = a.shape
    a4D = a.reshape(height,m//height,width,n//width)
    return reduction_func(a4D,axis=(1,3))

Thus, to solve our specific case, it would be :

blockwise_reduction(Z, height=5, width=5)

and for a block-wise average computation, it would be -

blockwise_reduction(Z, height=5, width=5, reduction_func=np.mean)