计算array1第i行和array2第i列的乘积 - NumPy

Question

I have a matrix M1 of shape (N*2) and another matrix M2 (2*N) , I want to obtain a result of (N) , each element i is the product of i th row of M1 and i th column of M2 . I tried to use dot in NumPy, but it can only give me the matrix multiplication result, which is (N*N) , of course, I can take the diagonal which is what I want, I would like to know is there a better way to do this?

Answer 1

Approach #1

You can use np.einsum -

np.einsum('ij,ji->i',M1,M2)

Explanation :

The original loopy solution would look something like this -

def original_app(M1,M2):
    N = M1.shape[0]
    out = np.zeros(N)
    for i in range(N):
        out[i] = M1[i].dot(M2[:,i])
    return out

Thus, for each iteration, we have :

out[i] = M1[i].dot(M2[:,i])

Looking at the iterator, we need to align the first axis of M1 with the second axis of M2 . Again, since we are performing matrix-multiplication and that by its very definition is aligning the second axis of M1 with the first axis of M2 and also sum-reducing these elements at each iteration.

When porting over to einsum , keep the axes to be aligned between the two inputs to have the same string when specifying the string notation to it. So, the inputs would be 'ij,ji for M1 and M2 respectively. The output after losing the second string from M1 , which is same as first string from M2 in that sum-reduction, should be left as i . Thus, the complete string notation would be : 'ij,ji->i' and the final solution as : np.einsum('ij,ji->i',M1,M2) .

Approach #2

The number of cols in M1 or number of rows in M2 is 2 . So, alternatively, we can just slice, perform the element-wise multiplication and sum up those, like so -

M1[:,0]*M2[0] + M1[:,1]*M2[1]

---

Runtime test

In [431]: # Setup inputs
     ...: N = 1000
     ...: M1 = np.random.rand(N,2)
     ...: M2 = np.random.rand(2,N)
     ...: 

In [432]: np.allclose(original_app(M1,M2),np.einsum('ij,ji->i',M1,M2))
Out[432]: True

In [433]: np.allclose(original_app(M1,M2),M1[:,0]*M2[0] + M1[:,1]*M2[1])
Out[433]: True

In [434]: %timeit original_app(M1,M2)
100 loops, best of 3: 2.09 ms per loop

In [435]: %timeit np.einsum('ij,ji->i',M1,M2)
100000 loops, best of 3: 13 µs per loop

In [436]: %timeit M1[:,0]*M2[0] + M1[:,1]*M2[1]
100000 loops, best of 3: 14.2 µs per loop

Massive speedup there!