Best way to take mean/sum of block matrix in numpy?

Question

I would like to perform some simple calculation on the block matrix (or more generally on the d-dim nd.array). Something like this:

In the picture, capital letters represent a 3 by 3 block matrix, and lower case letters are numbers (mean or sum of block matrices).

Currently, I only know how to do this by using for loop

import numpy as np

test_matrix = np.arange(81).reshape(9,9)
a = np.zeros((3,3))
for i in range(3):
    for j in range(3):
        a[k,i,j] = test_matrix[3*i:3*(i+1),3*j:3*(j+1)].mean()
print a

But it gets slow if my matrix gets bigger or multi-dimensional, such as the following:

test_matrix = np.arange(81*2).reshape(2,9,9)

a = np.zeros((2,3,3))
for k in range(2):
    for i in range(3):
        for j in range(3):
            a[k,i,j] = test_matrix[k,3*i:3*(i+1),3*j:3*(j+1)].mean()
print a

Is there any better way to perform this kind of tasks?

Thanks a lot!!

Answer 1

In [1952]: test=np.arange(81).reshape(9,9)
In [1953]: res=np.zeros((3,3))
In [1954]: for i in range(3):
      ...:     for j in range(3):
      ...:         res[i,j]=test[3*i:3*(i+1),3*j:3*(j+1)].mean()
In [1955]: res
Out[1955]: 
array([[ 10.,  13.,  16.],
       [ 37.,  40.,  43.],
       [ 64.,  67.,  70.]])

With reshape and sum or mean on selected axes:

In [1956]: test.reshape(3,3,3,3).mean(axis=(1,3))
Out[1956]: 
array([[ 10.,  13.,  16.],
       [ 37.,  40.,  43.],
       [ 64.,  67.,  70.]])

sum/mean allows us to specify 2 or more axes at a time, but it can also be done with repeated single axis applications.

test.reshape(3,3,3,3).mean(3).mean(1)

For the 3d array, these work

test.reshape( 2,3,3,3,3).mean((2,4))
test.reshape(-1,3,3,3,3).mean((2,4))

The -1 preserves the original 1st dimension (or in the case of a 2d test, it adds a size 1 dimension).