Summing over all but one axis

Question

Several times, I've dealt with an ND array, such as

foo = np.arange(27).reshape((3,3, 3))

and then I would have a dimension over which I want to preserve the variable from the next operation. Say the next operation is mean, in which case

preserveAxis = 1
desiredOutcome = foo.mean(axis=0).mean(axis=1)

the previous is my desired outcome, since I first take the average over the 0th axis, and then over the 2nd axis (which has become the 1st after the initial operation). That is, I've done the operation over axis 0 and 2, but preserved axis 1.

This type of procedure is cumbersome, and most importantly not generic. I'm looking for a generic way to preserve one axis, but sum/mean over all others. How could I achieve this best, preferably within numpy?

Answer 1

Here's one generalized to n-dim cases for reduction ufuncs -

def reduce_skipfew(ufunc, foo, preserveAxis=None):
    r = np.arange(foo.ndim)   
    if preserveAxis is not None:
        preserveAxis = tuple(np.delete(r, preserveAxis))
    return ufunc(foo, axis=preserveAxis)

Sample run -

In [171]: reduce_skipfew(np.mean, foo, preserveAxis=1)
Out[171]: array([10., 13., 16.])

In [172]: foo = np.arange(27).reshape((3,3, 3))

In [173]: reduce_skipfew(np.mean, foo, preserveAxis=1)
Out[173]: array([10., 13., 16.])

In [174]: reduce_skipfew(np.sum, foo, preserveAxis=1)
Out[174]: array([ 90, 117, 144])

# preserve none i.e. sum all
In [175]: reduce_skipfew(np.sum, foo, preserveAxis=None)
Out[175]: 351

# preserve more than one axis
In [176]: reduce_skipfew(np.sum, foo, preserveAxis=(0,2)) 
Out[176]: 
array([[ 9, 12, 15],
       [36, 39, 42],
       [63, 66, 69]])

Answer 2

You can use tuples for the axis parameter:

foo.mean(axis=(0, 2))

If your array have a variable number of dimensions and/or the preserved dimension is can vary, it can be al little more tricky. See @Divakar's answer.