Iterate in submatrices through a bigger matrix

Question

I looking for a way in which I using a for loop can iterate through a bigger matrix, in which each iteration will output a sub matrix of size (row, col, depth) (6,3,3).

My big matrix is stored as numpy matrix, and would possible also like the each iteration to be outputted as such.

>>> import numpy as np
>>> a = np.random.rand(6*3,3*3,3)
>>> print a.shape
(18, 9, 3)
>>> print a
>>> b

The variable b should contain all the sub matrixes of size (6,3,3) from matrix a. Each submatrix should not overlap with the prior.

Answer 1

Approach #1

I am assuming we are looking for non-overlapping/distinct blocks. As such we could use Scikit-image's view_as_blocks utility -

from skimage.util.shape import view_as_blocks

BSZ = (6,3,3)
out = view_as_blocks(a,BSZ).reshape((-1,)+ (BSZ))

Sample run -

In [279]: a = np.random.rand(6*3,3*3,3)

In [280]: out = view_as_blocks(a,BSZ).reshape((-1,)+ (BSZ))

In [281]: out.shape
Out[281]: (9, 6, 3, 3)

---

Approach #2

Using just native NumPy tools like reshaping and transpose, here's one way -

m,n,r = a.shape
split_shp = m//BSZ[0], BSZ[0], n//BSZ[1], BSZ[1], r//BSZ[2], BSZ[2]
out = a.reshape(split_shp).transpose(0,2,4,1,3,5).reshape((-1,)+ (BSZ))