Numpy Matrix Multiplication Broadcast

Question

I have an array comprised of N 3x3 arrays (a collection of matrices, although the data type is np.ndarray) and I have an array comprised of N 3x1 arrays (a collection of vectors). What I want to do is multiply each matrix by each vector, so I expect to get back N 3x1 arrays.

Simple example:

A = np.ones((6,3,3))
B = np.ones((6,3,1))
np.dot(A,B) # This gives me a 6x3x6x1 array, which is not what I want
np.array(map(np.dot,A,B)) # This gives me exactly what I want, but I don't want to have to rely on map

I've tired all kinds of reshaping, explored einsum, etc., but can't get this to work the way I want it to. How do I get this to work with numpy broadcasting? This operation will ultimately need to be done many thousands of times, and I don't want map or list comprehension operations to slow things down.

Answer 1

You can use np.einsum to calculate the dot products and create the matrix of the desired shape:

np.einsum('ijk,ikl->ijl', A, B)

Answer 2

A = np.random.rand(6, 3, 3)
B = np.random.rand(6, 3, 1)
C = np.array(map(np.dot, A, B))
D = np.sum(A*B.swapaxes(1, 2), axis=2)[..., None]
assert np.allclose(C, D)
assert C.shape == D.shape == (6, 3, 1)

The "allclose" is because there's some floating point rounding difference between the two methods on the order of 1e-16.

The .swapaxis and the [..., None] are just to get the arrays to conform to the shapes you specified. You could also represent it more simply with:

A = np.random.rand(6, 3, 3)
B = np.random.rand(6, 3)
C = np.array(map(np.dot, A, B))
D = np.sum(A*B[:, None, :], axis=2)
assert np.allclose(C, D)
assert C.shape == D.shape == (6, 3)

Answer 3

One could use the built-in matrix multiplication in Python 3.5 or above, introduced in PEP 465.

$ python --version
Python 3.6.6

>>> import numpy as np
>>> A = np.ones((6,3,3))
>>> B = np.ones((6,3,1))
>>> C = A @ B
>>> print(C)
[[[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]

 [[3.]
  [3.]
  [3.]]]