问题:

I have a numpy array with a shape of:

(11L, 5L, 5L)

I want to calculate the mean over the 25 elements of each 'slice' of the array [0, :, :], [1, :, :] etc, returning 11 values.

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

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

回答1:

Use a tuple for 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 works only with numpy version 1.7+.

回答2:

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 higher dimensional arrays (all of these methods would 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 boot:

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 then the other methods as array size increases.

Using dtype=np.float64 does not change the above timings appreciably, so just to 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 something that is 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