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

Posted 2023-02-23

【中文标题】在numpy中获取3D数组的2D切片的平均值

【英文标题】：Get mean of 2D slice of a 3D array in numpy

【发布时间】：2013-08-23 19:17:19

【问题描述】：

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

(11L, 5L, 5L)

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

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

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

【参考方案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+。

【讨论】：

有效吗？对于 1.7 及之后的版本，人们会这么认为，但文档仍然说只有一个轴。

没想到 numpy 版本，我有 1.7.1 并且可以工作。它不在文档中，但更新日志正在谈论 ufunc：softpedia.com/progChangelog/Numpy-Changelog-103892.html

酷，不知道这已经被添加了！

【参考方案2】：

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

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

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

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

【讨论】：

我认为这是最直接的答案，尽管 einsum 似乎确实更快。

【参考方案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

【讨论】：

我认为速度来自您使用int 累加器调用np.einsum，而不是float 或double，不确定np.mean 使用。这对于计算统计数据来说是一件冒险的事情，因为您可能会溢出累加器并得到非常错误的结果。给np.einsum 一个dtype=np.float 或dtype=np.double 都会使计算更加健壮，并且（我在这里猜测）在性能上与标准函数更相似。但是np.einsum 仍然是一个超酷的功能，所以你得到你的 +1...

@Jamie。这也是我的想法，但在我的初步测试中，einsum 实际上对于任何大小和 dtype 都更快。我已经用np.double时间更新了帖子。

@Ophion...很奇怪sum() 的速度与einsum() 的速度不同...很好观察...实际上第二种更快的计算平均值的方法是：timeit a.sum(axis=(1,2))/a.shape[-1]/a.shape[-2]

@Ophion 我认为您应该发布一个问题，例如“为什么np.einsum() 比np.sum() 快？”打开这个话题进行更详细的讨论...

@SaulloCastro 我刚才正在写一个这样的问题。使用a.sum(axis=(1,2)... 在时间上等同于a.mean(axis=(1,2)) 函数。