how to 3-way outer product in numpy?

Question

About the numpy.outer [link] .

Given two vectors, a = [a0, a1, ..., aM] and b = [b0, b1, ..., bN], the outer product will be M*N matrix.

Answer 1

You can use numpy.ix_; it adds axes to its operands such that they form an open grid. Below the shapes of A,B,C are (2, 1, 1), (1, 3, 1) and (1, 1, 4), so just multiplying them together results in the outer product.

a = np.arange(1, 3)
b = np.arange(1, 4)
c = np.arange(1, 5)
A,B,C = np.ix_(a,b,c)
A*B*C
# array([[[ 1,  2,  3,  4],
#         [ 2,  4,  6,  8],
#         [ 3,  6,  9, 12]],
#
#        [[ 2,  4,  6,  8],
#         [ 4,  8, 12, 16],
#         [ 6, 12, 18, 24]]])

Answer 2

You could use einsum for this:

>>> numpy.einsum('i,j,k->ijk', [1, 2], [3, 4, 5], [6,7,8,9])
array([[[18, 21, 24, 27],
        [24, 28, 32, 36],
        [30, 35, 40, 45]],

       [[36, 42, 48, 54],
        [48, 56, 64, 72],
        [60, 70, 80, 90]]])

---

You cannot cross arbitrary amount of vectors using einsum. The indices can only be letters, so at most 52 vectors are allowed in this form (although it will raise "too many operands" when 32 vectors are provided):

def nary_outer_einsum_52(*vectors):
    import string
    subscripts = string.ascii_letters[:len(vectors)]
    subscripts = ','.join(subscripts) + '->' + subscripts
    # expands to `numpy.einsum('a,b,c,d,e->abcde', v[0], v[1], v[2], v[3], v[4])`
    return numpy.einsum(subscripts, *vectors)

einsum supports another form which uses numeric indices instead of letters, which unfortunately only support up to 26 vectors because it just translates to the letter version internally. I do not recommend using this until the bug I mentioned is fixed.

def nary_outer_einsum_26(*vectors):
    operations = (x for i, v in enumerate(vectors) for x in (v, [i]))
    # expands to `numpy.einsum(v[0], [0], v[1], [1], v[2], [2], v[3], [3], v[4], [4])`
    return numpy.einsum(*operations)

numpy.ix_ (@PaulPanzer's answer) can easily be scaled up to the N-ary case. Still, there is an implementation-limit of 32 vectors:

def nary_outer_ix(*vectors):
    return numpy.prod(numpy.ix_(*vectors), axis=0)

---

timeit with 3 vectors:

>>> timeit.Timer('nary_outer_einsum_52([1,2], [3,4,5], [6,7,8,9])', globals=globals()).autorange()
(100000, 0.8991979530110257)   # 9 µs / iteration
>>> timeit.Timer('nary_outer_einsum_26([1,2], [3,4,5], [6,7,8,9])', globals=globals()).autorange()
(100000, 1.0089023629989242)   # 10 µs / iteration
>>> timeit.Timer('nary_outer_ix([1,2], [3,4,5], [6,7,8,9])', globals=globals()).autorange()
(10000, 0.30978472300921567)   # 31 µs / iteration

with 26 vectors:

>>> timeit.Timer('nary_outer_einsum_52(*([[1]] * 26))', globals=globals()).autorange()
(10000, 0.6589978060073918)   # 66 µs / iteration
>>> timeit.Timer('nary_outer_einsum_26(*([[1]] * 26))', globals=globals()).autorange()
(10000, 0.7502327310066903)   # 75 µs / iteration
>>> timeit.Timer('nary_outer_ix(*([[1]] * 26))', globals=globals()).autorange()
(1000, 0.2026848039968172)    # 203 µs / iteration

with 33 vectors:

>>> nary_outer_einsum_52(*([[1]] * 33))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 6, in nary_outer_einsum_52
  File "/usr/local/lib/python3.6/site-packages/numpy/core/einsumfunc.py", line 948, in einsum
    return c_einsum(*operands, **kwargs)
ValueError: too many operands
>>> nary_outer_ix(*([[1]] * 33))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in nary_outer_ix
  File "/usr/local/lib/python3.6/site-packages/numpy/lib/index_tricks.py", line 83, in ix_
    new = new.reshape((1,)*k + (new.size,) + (1,)*(nd-k-1))
ValueError: sequence too large; cannot be greater than 32

Answer 3

The direct way of doing this, taking full advantage of broadcasting is:

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

np.ix_ does this reshaping for you, at a modest cost in speed

In [919]: np.ix_(a,b,c)
Out[919]: 
(array([[[0]],

        [[1]],

        [[2]],

        [[3]],

        [[4]]]), array([[[10],
         [11],
         [12],
         [13]]]), array([[[20, 21, 22]]]))

and the resulting arrays can be multiplied with

np.prod(np.ix_(a,b,c))

The einsum version is simple, and fast

np.einsum('i,j,k',a,b,c)

It's a good idea to learn all 3 methods.

The problem with nesting outer is that expects the inputs to be 1d, or it flattens them. It can be used, but needs some reshaping

np.outer(a,np.outer(b,c)).reshape(a.shape[0],b.shape[0],c.shape[0])