numpy：沿除最后一个轴以外的所有轴求和

Question

如果我有一个任意形状的 ndarray，并且我想计算除最后一个轴以外的所有轴上的和，例如，我可以通过执行以下操作来实现：

all_but_last = tuple(range(arr.ndim - 1))
sum = arr.sum(axis=all_but_last)

现在，我感觉tuple(range(arr.ndim - 1))并不是很直观。有没有更优雅/类似numpy的方法来做到这一点？

此外，如果我想对多个形状不同的数组执行此操作，我需要为每个数组计算单独的维度元组。是否有更规范的方法来说“无论维度是什么，只给我除了一个轴以外的所有轴”？

- DarthPumpkin

---

直观与否，我认为这是最好利用“axis”参数的方法。将其隐藏在函数中，没有人会知道区别（您可以进行注释）。numpy函数经常玩这种“游戏”。 - hpaulj

Answer 1

你可以重塑数组，使除了最后一个轴之外的所有轴都被展平（例如，形状为(k, l, m, n)的数组变成(k*l*m, n)），然后在第一个轴上求和。

例如，以下是你的计算：

In [170]: arr.shape
Out[170]: (2, 3, 4)

In [171]: arr.sum(axis=tuple(range(arr.ndim - 1)))
Out[171]: array([2.85994792, 2.8922732 , 2.29051163, 2.77275709])

这是另一种选择：

In [172]: arr.reshape(-1, arr.shape[-1]).sum(axis=0)
Out[172]: array([2.85994792, 2.8922732 , 2.29051163, 2.77275709])

- Warren Weckesser

---

有没有更通用的方法？我想你可以做到这一点，这正是我正在寻找的。 sums = arr.sum(axis=tuple(x for x in range(arr.ndim) if x!= MY_AXIS)) 然后 sums.shape 就只会是 (MY_AXIS,)。 - momo668

Answer 2

您可以使用np.apply_over_axes对多个轴进行求和。

np.apply_over_axes(np.sum, arr, [0,2]) #sum over axes 0 and 2

np.apply_over_axes(np.sum, arr, range(arr.ndim - 1)) #sum over all but last axis

- Easton Bornemeier

---

这不就是做同样的事情吗？只是以更慢的方式，通过在Python代码中迭代轴而不是在“sum”内部进行操作？ - hpaulj

@hpaulj 我不确定你在Python代码中迭代所有轴的意思，因为这里展示的所有方法都使用range进行轴外推。 - Easton Bornemeier

但是看一下 apply_over_axes 的代码 - 它在 axes 参数上进行迭代。还有 timeit。 - hpaulj

抛开速度不谈，相比于sum(axis=...)，这种方法有什么好处？ - DarthPumpkin

@hpaulj 上述内容很难概括到任何轴上，因此为了快速开发，我使用了运行较慢的方法。 - Gulzar