如何计算 2D numpy 数组的所有列的总和（有效）

Posted 2023-03-12

【中文标题】如何计算 2D numpy 数组的所有列的总和（有效）

【英文标题】：How to calculate the sum of all columns of a 2D numpy array (efficiently)

【发布时间】：2012-11-14 01:51:07

【问题描述】：

假设我有以下由四行三列组成的二维 numpy 数组：

>>> a = numpy.arange(12).reshape(4,3)
>>> print(a)
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

生成包含所有列之和的一维数组（如[18, 22, 26]）的有效方法是什么？这可以在不需要遍历所有列的情况下完成吗？

【参考方案1】：

查看numpy.sum 的文档，特别注意axis 参数。总结列：

>>> import numpy as np
>>> a = np.arange(12).reshape(4,3)
>>> a.sum(axis=0)
array([18, 22, 26])

或者，对行求和：

>>> a.sum(axis=1)
array([ 3, 12, 21, 30])

其他聚合函数，例如numpy.mean、numpy.cumsum 和numpy.std，也采用axis 参数。

来自Tentative Numpy Tutorial：

许多一元运算，例如计算所有元素的总和 在数组中，被实现为ndarray 类的方法。经过 默认情况下，这些操作适用于数组，就好像它是一个列表一样 数字，不管它的形状。但是，通过指定 axis 参数，您可以沿指定轴应用操作 数组：

【讨论】：

抱歉，我不确定你的意思。使用 sum 函数对 numpy 数组的一个或多个轴求和。那是问题吗？你有别的想法吗？

这是一个很好的答案。不过，我通常更喜欢a.sum(axis=0) 而不是a.sum(0)。 （我认为它更明确一点——这绝不是一件坏事）

@Puggie，也许“更通用”是指“不使用内置 NumPy 函数”？一般来说，使用 NumPy 内置的函数会好得多，原因如下：它们已由 NumPy 开发团队优化，需要维护的代码更少，并且代码的可读性更高。 np.sum 函数在某种意义上是最通用的和最有效的，因为它隐藏了实现并且可能利用了 numpy 开发人员对 numpy 内部结构的了解。函数很好——使用它们。

@Puggie，啊，现在我明白你的意思了，尽管问题确实要求总和。在这种情况下，请参阅 np.apply_along_axis 和 np.apply_over_axes。

@JohnVinyard - 如果我只想对列（或行）的一个子集求和怎么办？有没有办法指定一组索引沿某个轴求和？谢谢！

【参考方案2】：

对列求和的其他替代方法是

numpy.einsum('ij->j', a)

和

numpy.dot(a.T, numpy.ones(a.shape[0]))

如果行数和列数在同一个数量级，那么所有的可能性都大致相同：

但是，如果只有几列，einsum 和 dot 解决方案的性能都明显优于 numpy 的 sum（注意对数刻度）：

---

重现图的代码：

import numpy
import perfplot


def numpy_sum(a):
    return numpy.sum(a, axis=1)


def einsum(a):
    return numpy.einsum('ij->i', a)


def dot_ones(a):
    return numpy.dot(a, numpy.ones(a.shape[1]))


perfplot.save(
    "out1.png",
    # setup=lambda n: numpy.random.rand(n, n),
    setup=lambda n: numpy.random.rand(n, 3),
    n_range=[2**k for k in range(15)],
    kernels=[numpy_sum, einsum, dot_ones],
    logx=True,
    logy=True,
    xlabel='len(a)',
    )

【参考方案3】：

使用axis 参数：

>> numpy.sum(a, axis=0)
  array([18, 22, 26])

【参考方案4】：

使用numpy.sum。对于你的情况，它是

sum = a.sum(axis=0)

【参考方案5】：

然后 NumPy sum 函数接受一个可选的轴参数，该参数指定您希望沿哪个轴执行总和：

>>> a = numpy.arange(12).reshape(4,3)
>>> a.sum(0)
array([18, 22, 26])

或者，等效地：

>>> numpy.sum(a, 0)
array([18, 22, 26])

【参考方案6】：

a.sum(0)

应该可以解决问题。这是一个 2d np.array，您将获得所有列的总和。 axis=0 是指向下方的维度，axis=1 是指向右侧的维度。