How to do an outer product of 3 vectors to create a 3d matrix in numpy? (and same for nd)

Viewed 295

If i want to do an outer product of 2 vectors to create a 2d matrix, each element a product of the two respective elements in the original vectors:

b = np.arange(5).reshape((1, 5))
a = np.arange(5).reshape((5, 1))
a * b
array([[ 0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16]])

I want the same for 3 (or for n) vectors.

An equivalent non numpy answer:

a = np.arange(5)
b = np.arange(5)
c = np.arange(5)
res = np.zeros((a.shape[0], b.shape[0], c.shape[0]))

for ia in range(len(a)):
    for ib in range(len(b)):
        for ic in range(len(c)):
            res[ia, ib, ic] = a[ia] * b[ib] * c[ic]
print(res)

out:

[[[ 0.  0.  0.  0.  0.]
  [ 0.  0.  0.  0.  0.]
  [ 0.  0.  0.  0.  0.]
  [ 0.  0.  0.  0.  0.]
  [ 0.  0.  0.  0.  0.]]

 [[ 0.  0.  0.  0.  0.]
  [ 0.  1.  2.  3.  4.]
  [ 0.  2.  4.  6.  8.]
  [ 0.  3.  6.  9. 12.]
  [ 0.  4.  8. 12. 16.]]

 [[ 0.  0.  0.  0.  0.]
  [ 0.  2.  4.  6.  8.]
  [ 0.  4.  8. 12. 16.]
  [ 0.  6. 12. 18. 24.]
  [ 0.  8. 16. 24. 32.]]

 [[ 0.  0.  0.  0.  0.]
  [ 0.  3.  6.  9. 12.]
  [ 0.  6. 12. 18. 24.]
  [ 0.  9. 18. 27. 36.]
  [ 0. 12. 24. 36. 48.]]

 [[ 0.  0.  0.  0.  0.]
  [ 0.  4.  8. 12. 16.]
  [ 0.  8. 16. 24. 32.]
  [ 0. 12. 24. 36. 48.]
  [ 0. 16. 32. 48. 64.]]]

How to do this with numpy [no for loops]?


Also, how to do this for a general function, not necessarily *?

3 Answers

NumPy provides you with np.outer() for computing the outer product. This is a less powerful version of more versatile approaches:

np.einsum() is the only one capable of handling more than two input arrays:

import numpy as np


def prod(items, start=1):
    for item in items:
        start = start * item
    return start


a = np.arange(5)
b = np.arange(5)
c = np.arange(5)


r0 = np.zeros((a.shape[0], b.shape[0], c.shape[0]))
for ia in range(len(a)):
    for ib in range(len(b)):
        for ic in range(len(c)):
            r0[ia, ib, ic] = a[ia] * b[ib] * c[ic]

r1 = prod([a[:, None, None], b[None, :, None], c[None, None, :]])
# same as: r1 = a[:, None, None] * b[None, :, None] * c[None, None, :]
# same as: r1 = a.reshape(-1, 1, 1) * b.reshape(1, -1, 1) * c.reshape(1, 1, -1)
print(np.all(r0 == r2))
# True

r2 = np.einsum('i,j,k->ijk', a, b, c)
print(np.all(r0 == r2))
# True

# as per @hpaulj suggestion
r3 = prod(np.ix_(a, b, c))
print(np.all(r0 == r3))
# True

Of course, the broadcasting approach (which is the same that you used with the array.reshape() version of your code, except that it uses a slightly different syntax for providing the correct shape), can be automatized by explicitly building the slicing (or equivalently the array.reshape() parameters).

In [166]: a = np.arange(2)
     ...: b = np.arange(3)
     ...: c = np.arange(4)

As shown in comments and answer:

In [167]: R = np.einsum('i,j,k',a,b,c)

We can also np.ix_ construct arrays that broadcast against each other. This is often used to construct block indexing arrays, but works here as well:

In [168]: A,B,C = np.ix_(a,b,c)
In [169]: A,B,C
Out[169]: 
(array([[[0]],
 
        [[1]]]),
 array([[[0],
         [1],
         [2]]]),
 array([[[0, 1, 2, 3]]]))
In [170]: R1 = A*B*C

testing:

In [171]: np.allclose(R,R1)
Out[171]: True

That broadcasted product can be done in one line with:

In [172]: np.prod(np.array(np.ix_(a,b,c),object)).shape
Out[172]: (2, 3, 4)

Without that explicit object dtype casting I get a future warning about creating an ragged array.

np.meshgrid(a,b,c, sparse=True, indexing='ij') is an alternative to ix_.

While these ix_ etc expressions are nice, you should become thoroughly comfortable using:

a[:, None, None] * b[None, :, None] * c[None, None, :]

This kind of dimension expansion gives you the most power and flexibility.

The simplest approach (from this answer) is to use:

functools.reduce(np.multiply.outer, (a, b, c))

This works for any number of dimensions, and unlike np.prod(np.ix_(...)) it does not result in numpy deprecation warnings about introducing jagged arrays.

Related