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