CODEWITHSUNDEEP
pythonarraysnumpymatrix
Hendrik Wiese
Hendrik Wiese

Reputation: 2219

Multiply array of vectors with array of matrices; return array of vectors?

I've got a numpy array of row vectors of shape (n,3) and another numpy array of matrices of shape (n,3,3). I would like to multiply each of the n vectors with the corresponding matrix and return an array of shape (n,3) of the resulting vectors.

By now I've been using a for loop to iterate through the n vectors/matrices and do the multiplication item by item.

I would like to know if there's a more numpy-ish way of doing this. A way without the for loop that might even be faster.

//edit 1:

As requested, here's my loopy code (with n = 10):

    arr_in = np.random.randn(10, 3)
    matrices = np.random.randn(10, 3, 3)

    for i in range(arr_in.shape[0]): # 10 iterations
        arr_out[i] = np.asarray(np.dot(arr_in[i], matrices[i]))

Upvotes: 5

Views: 1340

Answers (2)

Divakar
Divakar

Reputation: 221684

That dot-product is essentially performing reduction along axis=1 of the two input arrays. The dimensions could be represented like so -

arr_in   :     n   3 
matrices : n   3   3

So, one way to solve it would be to "push" the dimensions of arr_in to front by one axis/dimension, thus creating a singleton dimension at axis=2 in a 3D array version of it. Then, sum-reducing the elements along axis = 1 would give us the desired output. Let's show it -

arr_in   : n   [3]   1 
matrices : n   [3]   3

Now, this could be achieved through two ways.

1) With np.einsum -

np.einsum('ij,ijk->ik',arr_in,matrices)

2) With NumPy broadcasting -

(arr_in[...,None]*matrices).sum(1)

Runtime test and verify output (for einsum version) -

In [329]: def loop_based(arr_in,matrices):
     ...:     arr_out = np.zeros((arr_in.shape[0], 3))
     ...:     for i in range(arr_in.shape[0]):
     ...:         arr_out[i] =  np.dot(arr_in[i], matrices[i])
     ...:     return arr_out
     ...: 
     ...: def einsum_based(arr_in,matrices):
     ...:     return np.einsum('ij,ijk->ik',arr_in,matrices)
     ...: 

In [330]: # Inputs
     ...: N = 16935
     ...: arr_in = np.random.randn(N, 3)
     ...: matrices = np.random.randn(N, 3, 3)
     ...: 

In [331]: np.allclose(einsum_based(arr_in,matrices),loop_based(arr_in,matrices))
Out[331]: True

In [332]: %timeit loop_based(arr_in,matrices)
10 loops, best of 3: 49.1 ms per loop

In [333]: %timeit einsum_based(arr_in,matrices)
1000 loops, best of 3: 714 µs per loop

Upvotes: 2

Bill
Bill

Reputation: 488

You could use np.einsum. To get v.dot(M) for each vector-matrix pair, use np.einsum("...i,...ij", arr_in, matrices). To get M.dot(v) use np.einsum("...ij,...i", matrices, arr_in)

Upvotes: 0

Related Questions