Blockwise operations in Numpy

Question

Are there any convenience utilities for doing blockwise operations on Numpy arrays?

I am thinking of operations like Ising spin renormalization where you divide a matrix into blocks and return matrix where each block is replaced by its sum, average or other function.

Answer 1

You might be looking for superbatfish's blockwise_view. This uses np.lib.stride_tricks.as_strided to create a view of the array which places "blocks" of the array in their own axes.

For example, suppose you have a 2D array such as,

In [97]: arr = np.arange(24).reshape(6, 4)

In [98]: arr.shape
Out[98]: (6, 4)

In [99]: arr
Out[99]: 
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]])

and you wish to "chop it" into 4 blocks of shape (3, 2). You could use blockwise_view to convert it into a 4D array of shape (4, 3, 2):

In [34]: blocked = blockwise_view(arr, (3, 2)); blocked
Out[34]: 
array([[[[ 0,  1],
         [ 4,  5],
         [ 8,  9]],

        [[ 2,  3],
         [ 6,  7],
         [10, 11]]],


       [[[12, 13],
         [16, 17],
         [20, 21]],

        [[14, 15],
         [18, 19],
         [22, 23]]]])

In [37]: blocked.shape
Out[37]: (2, 2, 3, 2)

Now you could reshape it so all the values from one block are in the last axis:

In [41]: reshaped = blocked.reshape(-1, 3*2); reshaped
Out[41]: 
array([[ 0,  1,  4,  5,  8,  9],
       [ 2,  3,  6,  7, 10, 11],
       [12, 13, 16, 17, 20, 21],
       [14, 15, 18, 19, 22, 23]])

Now you can sum along that axis, or take its mean or apply some other function to the elements of each block:

In [103]: reshaped.sum(axis=-1)
Out[103]: array([ 27,  39,  99, 111])

In [104]: reshaped.mean(axis=-1)
Out[104]: array([  4.5,   6.5,  16.5,  18.5])

Unlike my first answer, which can only be applied to 2D arrays, blockwise_view can be applied to arbitrary N-dimensional arrays. It returns a 2N-dimensional array where the first N axes index the blocks.