CODEWITHSUNDEEP
numpy
Make42
Make42

Reputation: 13088

Calculating the sum of vectors

I have two matrices a and b and would like to calculate all the sums between them into a tensor. How can I do this more efficiently than doing the following code:

a = np.array([[1,2],[3,4],[5,6]])
b = np.array([[4,5],[6,7]])

n1 = a.shape[0]
n2 = b.shape[0]
f = a.shape[1]
c = np.zeros((n1,n2,f))

c = np.zeros((n1,n2,f))
for i in range(n1):
    for j in range(n2):
        c[i,j,:] = a[i,:] + b[j,:]

einstein-sum and the like does obviously not work and an outer product neither - is there an appropriate method?

Upvotes: 0

Views: 60

Answers (2)

tstanisl
tstanisl

Reputation: 14127

Use broadcasting and add extra dimensions using advanced indexing:

a[:,None,:]+b[None,:,:]

Upvotes: 1

hpaulj
hpaulj

Reputation: 231395

You can transform your loop expression into a broadcasting one:

c[i,j,:] = a[i,:] + b[j,:]
c[i,j,:] = a[i,None,:] + b[None,j,:]  # fill in the missing dimensions
c = a[:,None,:] + b[None,:,:]


In [167]: a[:,None,:]+b[None,:,:]                                                                              
Out[167]: 
array([[[ 5,  7],
        [ 7,  9]],

       [[ 7,  9],
        [ 9, 11]],

       [[ 9, 11],
        [11, 13]]])
In [168]: _.shape                                                                                            
Out[168]: (3, 2, 2)

a[:,None]+b does the same thing, since leading None (np.newaxis) are automatic, and trailing : also.

Upvotes: 1

Related Questions