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

Question

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

arraysnumpymultidimensional-arrayslicemean

21

我有一个形状为：

(11L, 5L, 5L)

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

看起来很傻，但我无法弄清楚如何做到这一点。我曾尝试使用 mean(axis=x) 函数来实现，但我尝试了所有可能的轴组合，都没有得到我想要的结果。

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

- robintw

3个回答

8

您可以使用reshape(11, 25)来重新塑形数组，然后仅调用一次mean函数（速度更快）：

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

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

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

- Saullo G. P. Castro

1

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

5

您可以始终使用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

- Daniel

1

我觉得速度是因为你调用np.einsum时使用了一个int累加器，而不是float或double（不确定），而np.mean使用的是。在计算统计数据时，这样做是有风险的，因为你可能会溢出累加器并得到非常错误的结果。给np.einsum传递dtype=np.float或dtype=np.double将使计算更加稳健，并且（我猜测）性能与标准函数更相似。但是np.einsum仍然是一个超酷的函数，所以你得到了+1... - Jaime

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

@Ophion...很奇怪sum()不能像einsum()一样快...非常好的观察...实际上，计算平均值的第二种更快的方法是：timeit a.sum(axis=(1,2))/a.shape[-1]/a.shape[-2]。 - Saullo G. P. Castro

1

@Ophion 我认为您应该发一篇像“为什么np.einsum()比np.sum()更快？”这样的问题来深入讨论此话题... - Saullo G. P. Castro

1

@SaulloCastro 我刚刚写了一个类似的问题。使用 a.sum(axis=(1,2)... 与 a.mean(axis=(1,2)) 函数在时间上是等效的。 - Daniel

@Ophion 我刚看到了！很好的问题！ - Saullo G. P. Castro

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- J. Martinot-Lagarde · Accepted Answer

使用元组作为轴:

>>> 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及以上。