Replace looping-over-axes with broadcasting

Question

Say I have:

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

I want to get to:

c = np.array([[20,32],
              [28, 44]])

where c is the result of multiplying each column of a by b, then summing that result along the first axis.

I.e.:

print(np.sum(a[:, 0] * b, axis=1))
[20 32]

print(np.sum(a[:, 1] * b, axis=1))
[28 44]

Can I do through broadcasting rather than:

Divakar · Accepted Answer · 2017-07-19T20:47:19.780

4

You can use np.dot -

b.dot(a).T

Alternatively, using np.einsum (for the kicks maybe) -

np.einsum('ij,ki->jk',a,b)

edited Jul 19 '17 at 20:47

answered Jul 19 '17 at 20:44

Divakar

1

1+ for the kick. – cs95 Jul 19 '17 at 20:47
And `einsum` will be slower for smaller arrays and faster for larger ones? – Brad Solomon Jul 19 '17 at 20:51
@BradSolomon With no axis aligning requirement, I would think that as you increase the number of sum-reductions, i.e. work with large arrays, `np.dot` would fair better. – Divakar Jul 19 '17 at 20:54

1 Answers1