轴上的numpy sum

Question

是否有一个numpy的函数来汇总的阵列沿(没有结束)一个给定的轴线？沿着轴,我的意思是相当于:

[x.sum() for x in arr.swapaxes(0,i)].

沿轴i 求和.

例如,numpy.sum不能直接工作的情况:

>>> a = np.arange(12).reshape((3,2,2))
>>> a
array([[[ 0,  1],
        [ 2,  3]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 8,  9],
        [10, 11]]])
>>> [x.sum() for x in a] # sum along axis 0
[6, 22, 38]
>>> a.sum(axis=0)
array([[12, 15],
       [18, 21]])
>>> a.sum(axis=1)
array([[ 2,  4],
       [10, 12],
       [18, 20]])
>>> a.sum(axis=2)
array([[ 1,  5],
       [ 9, 13],
       [17, 21]])

Answer 1

您可以只传递一个包含要求和的轴的元组，并省略要“求和”的轴：

>> a.sum(axis=(1,2))
array([ 6, 22, 38])

Answer 2

def sum_along_axis(a, axis=None):
    """Equivalent to [x.sum() for x in a.swapaxes(0,axis)]"""
    if axis is None:
        return a.sum()
    return np.fromiter((x.sum() for x in a.swapaxes(0,axis)), dtype=a.dtype)