9

Say I have two arrays a and b,

  a.shape = (5,2,3)
  b.shape = (2,3)

then c = a * b will give me an array c of shape (5,2,3) with c[i,j,k] = a[i,j,k]*b[j,k].

Now the situation is,

  a.shape = (5,2,3)
  b.shape = (2,3,8)

and I want c to have a shape (5,2,3,8) with c[i,j,k,l] = a[i,j,k]*b[j,k,l].

How to do this efficiently? My a and b are actually quite large.

user1342516
  • 447
  • 2
  • 5
  • 10

2 Answers2

13

This should work:

a[..., numpy.newaxis] * b[numpy.newaxis, ...]

Usage:

In : a = numpy.random.randn(5,2,3)

In : b = numpy.random.randn(2,3,8)

In : c = a[..., numpy.newaxis]*b[numpy.newaxis, ...]

In : c.shape
Out: (5, 2, 3, 8)

Ref: Array Broadcasting in numpy

Edit: Updated reference URL

Avaris
  • 35,883
  • 7
  • 81
  • 72
  • I find that `None` also works, in addition to `numpy.newaxis`; and indeed, that `numpy.newaxis is None` is true for me. Is there a reason not to just use `None` here? – senderle May 20 '12 at 23:00
  • 1
    @senderle. Apparently, they are interchangable: [numpy.newaxis](http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#numpy.newaxis). But, in the future it might change. I guess that's why they put `newaxis` there as a synonym for `None`. – Avaris May 20 '12 at 23:06
  • `np.newaxis is None` returns `True`. I tend to use `newaxis` just to make things explicit in my code. Also see http://stackoverflow.com/questions/944863/numpy-should-i-use-newaxis-or-none – JoshAdel May 20 '12 at 23:07
  • @Avaris it will be grateful if u upvote my comment if it was really appreciable. – timekeeper Dec 13 '15 at 18:25
7

I think the following should work:

import numpy as np
a = np.random.normal(size=(5,2,3))
b = np.random.normal(size=(2,3,8))
c = np.einsum('ijk,jkl->ijkl',a,b)

and:

In [5]: c.shape
Out[5]: (5, 2, 3, 8)

In [6]: a[0,0,1]*b[0,1,2]
Out[6]: -0.041308376453821738

In [7]: c[0,0,1,2]
Out[7]: -0.041308376453821738

np.einsum can be a bit tricky to use, but is quite powerful for these sort of indexing problems:

http://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html

Also note that this requires numpy >= v1.6.0

I'm not sure about efficiency for your particular problem, but if it doesn't perform as well as needed, definitely look into using Cython with explicit for loops, and possibly parallelize it using prange

UPDATE

In [18]: %timeit np.einsum('ijk,jkl->ijkl',a,b)
100000 loops, best of 3: 4.78 us per loop

In [19]: %timeit a[..., np.newaxis]*b[np.newaxis, ...]
100000 loops, best of 3: 12.2 us per loop


In [20]: a = np.random.normal(size=(50,20,30))
In [21]: b = np.random.normal(size=(20,30,80))

In [22]: %timeit np.einsum('ijk,jkl->ijkl',a,b)
100 loops, best of 3: 16.6 ms per loop

In [23]: %timeit a[..., np.newaxis]*b[np.newaxis, ...]
100 loops, best of 3: 16.6 ms per loop

In [2]: a = np.random.normal(size=(500,20,30))
In [3]: b = np.random.normal(size=(20,30,800))

In [4]: %timeit np.einsum('ijk,jkl->ijkl',a,b)
1 loops, best of 3: 3.31 s per loop

In [5]: %timeit a[..., np.newaxis]*b[np.newaxis, ...]
1 loops, best of 3: 2.6 s per loop
JoshAdel
  • 66,734
  • 27
  • 141
  • 140
  • Umm, well... For larger arrays `numpy.einsum` seems to be a little bit slower. – Avaris May 20 '12 at 22:51
  • @Avaris I agree generally with your assessment of the timings, although over a pretty large range of array sizes, `np.einsum` is as fast or faster than the broadcasting solution. In some ways, I prefer the einsum solution for readability and making things explicit, although others might disagree. But certainly if for the use case, broadcasting is faster, use that. – JoshAdel May 20 '12 at 23:05