Search code examples
numpyvectormatrix-multiplication

Compact and natural way to write matrix product of vectors in Numpy


In scientific computing I often want to do vector multiplications like

a x b^T

with a and b being row vectors and b^T is the transpose of the vector. So if a and b are of shape [n, 1] and [m, 1], the resulting matrix has shape [n, m]

Is there a good and straight forward way to write this multiplication in numpy?

Example:

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

Adding axes manually works:

a[:,np.newaxis] @ b[np.newaxis,:]

and gives the correct result:

[[ 4  5  6  7]
 [ 8 10 12 14]
 [12 15 18 21]]

Einstein notation would be another way, but still somewhat weird.

np.einsum('a,b->ab', a,b)

What I was hoping to work, but doesn't work, is the following:

a @ b.T

Any other approaches to do this?


Solution

  • In MATLAB matrix multiplication is the norm, using *. Element wise multiplication uses .* operator. Also matrices are atleast 2d.

    In numpy, elementwise multiplication uses *. Matrix multiplication is done with np.dot (or its method), and more recently with the @ operator (np.matmul). numpy adds broadcasting, which gives the elementwise multiplication a lot more expresiveness.

    With your 2 examples arrays, of shape (3,) and (4,) the options of making a (3,4) outer product https://en.wikipedia.org/wiki/Outer_product include:

    np.outer(a,b)     
    
    np.einsum('i,j->ij, a, b)  # matching einstein index notation
    
    a[:,None] * b       # the most idiomatic numpy expression
    

    This last works because of broadcasting. a[:, None], like a.reshape(-1,1) turns the (3,) array into a (3,1). b[None, :] turns a (4,) into (1,4). But broadcasting can perform this upgrade automatically (and unambiguously).

    (3,1) * (4,) => (3,1) * (1,4) => (3,4)
    

    Broadcasting does not work with np.dot. So we need

    a[:, None].dot(b[None, :])   #  (3,1) dot with (1,4)
    

    The key with dot is that the last dim of a pairs with the 2nd to last of b. (np.dot also works with 2 matching 1d arrays, performing the conventional vector dot product).

    @ (matmul) introduces an operator that works like dot, at least in the 2d with 2d case. With higher dimensional arrays they work differently.

    a[:,None].dot(b[None,:])
    np.dot(a[:,None], b[None,:])
    a[:,None] @ b[None,:]
    a[:,None] @ b[:,None].T
    

    and the reshape equivalents all create the desired (3,4) array.

    np.tensordot can handle other dimensions combinations, but it works by reshaping and transposing the inputs, so in the end it can pass them to dot. It then transforms the result back into desired shape.

    Quick time tests show that np.dot versions tend to be fastest - because they delegate the action to fast BLAS like libraries. For the other versions, the delegation is a bit more indirect, or they use numpy's own compiled code.