CODEWITHSUNDEEP
pythonarraysnumpymultiplication
user1342516
user1342516

Reputation: 539

Elementwise multiplication of arrays of different shapes in python

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.

Upvotes: 9

Views: 4462

Answers (2)

Avaris
Avaris

Reputation: 36715

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

Upvotes: 13

JoshAdel
JoshAdel

Reputation: 68682

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

Upvotes: 7

Related Questions