2011-08-21 13 views
6

Esiste una funzione numpy per sommare un array lungo (non oltre) un dato asse? Con lungo un asse, voglio dire qualcosa di equivalente a:numpy sum lungo l'asse

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

Per riassumere lungo Asse I.

Per esempio, un caso in cui numpy.sum non funzionerà direttamente: somma

>>> 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]]) 

risposta

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) 
3

chiamare due volte?

In [1]: a.sum(axis=1).sum(axis=1) 
Out[1]: array([ 6, 22, 38]) 

Naturalmente, questo sarebbe un po 'scomodo da generalizzare perché gli assi "spariscono". Hai bisogno che sia generale?

def sum_along(a, axis=0): 
    js = [axis] + [i for i in range(len(a.shape)) if i != axis] 
    a = a.transpose(js) 

    while len(a.shape) > 1: a = a.sum(axis=1) 

    return a 
2
np.apply_over_axes(sum, a, [1,2]).ravel() 
3

A partire dal NumPy 1.7.1 c'è una risposta facile qui - è possibile passare una tupla all'argomento "asse" del metodo Somma per sommare su più assi. Quindi, per riassumere su tutto tranne quello dato:

x.sum(tuple(j for j in xrange(x.ndim) if j!=i)) 
2

si può semplicemente passare una tupla con gli assi che si desidera sommare sopra, e lasciare fuori quello che si desidera 'sum lungo':

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