Search code examples
pythonnumpynumpy-ufunc

Numpy power ufunc operating on specific axis


I find it weird that numpy.power has no axis argument... is it because there is a better/safer way to achieve the same goal (elevating each 2D array in a 3D array to the power of a 1D array).

Suppose you have a (3,10,10) array (A) and you want to elevate each (10,10) array to the power of elements in array B of shape (3,). You should be able to do it by using np.power(A,B,axis=0), right? Yet it yields the following TypeError :

TypeError: 'axis' is an invalid keyword to ufunc 'power'

Since it seems that power does not have an axis or axes argument (despite being an ufunc), what is the preferred way to do it ?

There may be a solution using the ufunc.reduce method but I don't really see how that would work with numpy.power...

For now I do :

np.array([A[i,:,:]**B[i] for i in range(3)])

But it looks ugly and is probably less efficient than a numpy method would be.

Thanks


Solution

  • power is not a reduction operation: it does not reduce a collection of numbers to a single number, so an axis argument doesn't make sense. Operations such as sum or max are reductions, so it is meaningful to specify an axis along which to apply the reduction.

    The operation that you want is broadcasting. Here's a smaller example, with A having shape (3, 2, 2) and B having shape (3,). We can't write np.power(A, B), because the shapes are not compatible for broadcasting. We first have to add trivial dimensions to B to give it the shape (3, 1, 1). That can be done with, for example, B[:, np.newaxis, np.newaxis] or B.reshape(-1, 1, 1).

    In [100]: A                                                                                                                                                    
    Out[100]: 
    array([[[1, 1],
            [3, 3]],
    
           [[3, 2],
            [1, 1]],
    
           [[3, 2],
            [1, 3]]])
    
    In [101]: B                                                                                                                                                    
    Out[101]: array([2, 1, 3])
    
    In [102]: np.power(A, B[:, np.newaxis, np.newaxis])                                                                                                            
    Out[102]: 
    array([[[ 1,  1],
            [ 9,  9]],
    
           [[ 3,  2],
            [ 1,  1]],
    
           [[27,  8],
            [ 1, 27]]])
    

    The value of np.newaxis is None, so you'll often see expressions that use None instead of np.newaxis. You can also using the ** operator instead of the function power:

    In [103]: A ** B[:, None, None]                                                                                                                                
    Out[103]: 
    array([[[ 1,  1],
            [ 9,  9]],
    
           [[ 3,  2],
            [ 1,  1]],
    
           [[27,  8],
            [ 1, 27]]])