Indices of fixed size sub-matrices of numpy array

Question

I am implementing an algorithm which requires me to look at non-overlapping consecutive submatrices within a (strictly two dimensional) numpy array. eg, for the 12 by 12

Answer 1

I'm adding this answer to an old question since an edit has bumped this question up. Here's an alternative way to calculate blocks:

size = 3
lenr, lenc = int(a.shape[0]/size), int(a.shape[1]/size)

t = a.reshape(lenr,size,lenc,size).transpose(0, 2, 1, 3)

Profiling shows that this is the fastest. Profiling done with python 3.5, and the results from map passed to array() for compatibility, since in 3.5 map returns an iterator.

reshape/transpose:   643 ns per loop
reshape/index:       45.8 µs per loop
Map/split:           10.3 µs per loop

It's interesting that the iterator version of map is faster. In any case, using reshape and transpose is fastest.

Answer 2

Getting the indexes

Here's a quick way to get a specific size x size block:

base = np.arange(size) # Just the base set of indexes
row = 1                # Which block you want
col = 0                
block = a[base[:, np.newaxis] + row * size, base + col * size]

If you wanted you could build up matrices similar to your xcoords like:

y, x = np.mgrid[0:a.shape[0]/size, 0:a.shape[1]/size]
y_coords = y[..., np.newaxis] * size + base
x_coords = x[..., np.newaxis] * size + base

Then you could access a block like this:

block = a[y_coords[row, col][:, np.newaxis], x_coords[row, col]]

---

Getting the blocks directly

If you just want to get the blocks (and not the indexes of the block entries), I'd use np.split (twice):

blocks = map(lambda x : np.split(x, a.shape[1]/size, 1), # Split the columns
                        np.split(a, a.shape[0]/size, 0)) # Split the rows

then you have a 2D list of size x size blocks:

>>> blocks[0][0]
array([[ 4,  0, 12],
       [15, 13,  2],
       [18, 18,  3]])

>>> blocks[1][0]
array([[ 1,  9,  3],
       [ 5, 15,  5],
       [13, 17,  8]])

You could then make this a numpy array and use the same indexing style as above:

>>> blocks = np.array(blocks)
>>> blocks.shape
(4, 4, 3, 3)

Answer 3

You can use the one-liner:

r = 3
c = 3
lenr = a.shape[0]/r
lenc = a.shape[1]/c
np.array([a[i*r:(i+1)*r,j*c:(j+1)*c] for (i,j) in np.ndindex(lenr,lenc)]).reshape(lenr,lenc,r,c)