在numpy中获取3D数组的2D切片的平均值

Question

我有一个 numpy 数组，其形状为：

(11L, 5L, 5L)

我想计算数组 [0, :, :], [1, :, :] 等的每个“切片”的 25 个元素的平均值，返回 11 个值。

这看起来很傻，但我不知道该怎么做。我原以为 mean(axis=x) 函数可以做到这一点，但我尝试了所有可能的轴组合，但都没有给我想要的结果。

我显然可以使用 for 循环和切片来做到这一点，但肯定有更好的方法吗？

Answer 1

对轴使用元组：

>>> 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.])

编辑：这只适用于 numpy 版本 1.7+。

Answer 2

您可以reshape(11, 25) 然后只调用一次mean（更快）：

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

或者，您可以调用np.mean 两次（在我的计算机上大约慢 2 倍）：

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

Answer 3

可以一直使用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])

适用于更高维的数组（如果轴标签发生更改，所有这些方法都可以）：

>>> 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)

启动速度更快：

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

随着数组大小的增加，比其他方法稍微好一些。

使用dtype=np.float64 不会明显改变上述时间，所以请仔细检查：

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

还有一些有趣的事情：

%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