Get average value of a 2D fragment of a 3D array in numpy

Question

I have a numpy array with the form:

(11L, 5L, 5L) 

I want to calculate the average of 25 elements of each “slice” of the array [0,:,:], [1,:,:], etc., returning 11 values.

Seems silly, but I can't figure out how to do this. I thought the mean(axis=x) function would do this, but I tried all possible combinations of the axis, and none of them gave me the result I want.

I obviously can do this using a for and slicing loop, but of course, is there a better way?

Answer 1

Use the tuple for the axis:

 >>> a = np.arange(11*5*5).reshape(11,5,5) >>> a.mean(axis=(1,2)) array([ 12., 37., 62., 87., 112., 137., 162., 187., 212., 237., 262.]) 

Edit: this only works with numpy version 1.7+.

Answer 2

You can always use np.einsum :

 >>> a = np.arange(11*5*5).reshape(11,5,5) >>> np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2]) array([ 12, 37, 62, 87, 112, 137, 162, 187, 212, 237, 262]) 

Works on arrays with large sizes (all these methods will be if the axis labels are changed):

 >>> a = np.arange(10*11*5*5).reshape(10,11,5,5) >>> (np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2])).shape (10, 11) 

Faster to load:

 a = np.arange(11*5*5).reshape(11,5,5) %timeit a.reshape(11, 25).mean(axis=1) 10000 loops, best of 3: 21.4 us per loop %timeit a.mean(axis=(1,2)) 10000 loops, best of 3: 19.4 us per loop %timeit np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2]) 100000 loops, best of 3: 8.26 us per loop 

Scales slightly better than other methods as the size of the array increases.

Using dtype=np.float64 does not change the above timings noticeably, so just double check:

 a = np.arange(110*50*50,dtype=np.float64).reshape(110,50,50) %timeit a.reshape(110,2500).mean(axis=1) 1000 loops, best of 3: 307 us per loop %timeit a.mean(axis=(1,2)) 1000 loops, best of 3: 308 us per loop %timeit np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2]) 10000 loops, best of 3: 145 us per loop 

Also interesting:

 %timeit np.sum(a) #37812362500.0 100000 loops, best of 3: 293 us per loop %timeit np.einsum('ijk->',a) #37812362500.0 100000 loops, best of 3: 144 us per loop 

Answer 3

You can reshape(11, 25) and then call mean only once (faster):

 a.reshape(11, 25).mean(axis=1) 

Alternatively, you can call np.mean twice (about 2 times slower on my computer):

 a.mean(axis=2).mean(axis=1)