Search code examples
pythonnumpysumaxisouter-join

Numpy: Outer sum along a specific axis


I'm looking for an efficient way to do an outer sum over a 2d matrix columns.

Let say A is a (2, 3) matrix, my wanted result is shaped (3,2,2) and defined as :

A = np.array([[a11, a12, a13],
              [a21, a22, a23]])

myWantedResult = np.array( [[[a11+a11, a11+a21],
                             [a21+a11, a21+a21]], # This is the np.add.outer(A[:,0],A[:,0])
                            [[a12+a12, a12+a22 ],
                             [a22+a12, a22+a22]], # This is the np.add.outer(A[:,1],A[:,1])
                           [[a13+a13, a13+a23],
                            [a23+a13, a23+a23]]# This is the np.add.outer(A[:,2],A[:,2])
                           ])

I've already tried to loop over the column but it is quite time-consuming because I want to do it on large array. I'm looking for a vectorized solution.

Many thanks !


Solution

  • The outer add can be done with broadcasting by adding A to a reshaped version of A with shape (2, 1, 3). That reshape can be done several ways, including A[:, None, :], A.reshape((2, 1, 3)) or np.expand_dims(A, 1). Then A + A[:, None, :] contains the result that you want, but the shape is (2, 2, 3). You can use np.moveaxis to rearrange the axes so that the shape is (3, 2, 2).

    In [174]: A                                                                                   
    Out[174]: 
    array([[  1,  10, 100],
           [-25,  50, 500]])
    
    In [175]: B = np.moveaxis(A + A[:, None, :], -1, 0)                                           
    
    In [176]: B.shape                                                                             
    Out[176]: (3, 2, 2)
    
    In [177]: B                                                                                   
    Out[177]: 
    array([[[   2,  -24],
            [ -24,  -50]],
    
           [[  20,   60],
            [  60,  100]],
    
           [[ 200,  600],
            [ 600, 1000]]])