Skip to content
Advertisement

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

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 *?

Advertisement

Answer

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).

User contributions licensed under: CC BY-SA
8 People found this is helpful
Advertisement