CODEWITHSUNDEEP
pythonnumpy
cp3028
cp3028

Reputation: 345

Numpy Array Broadcasting with different dimensions

I a little confused by the broadcasting rules of numpy. Suppose you want to perform an axis-wise scalar product of a higher dimension array to reduce the array dimension by one (basically to perform a weighted summation along one axis):

from numpy import *

A = ones((3,3,2))
v = array([1,2])

B = zeros((3,3))

# V01: this works
B[0,0] = v.dot(A[0,0])

# V02: this works
B[:,:] = v[0]*A[:,:,0] + v[1]*A[:,:,1] 

# V03: this doesn't
B[:,:] = v.dot(A[:,:])

Why does V03 not work?

Cheers

Upvotes: 7

Views: 3606

Answers (3)

Jaime
Jaime

Reputation: 67417

np.dot(a, b) operates over the last axis of a and the second-to-last of b. So for your particular case in your question,you could always go with :

>>> a.dot(v)
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

If you want to keep the v.dot(a) order, you need to get the axis into position, which can easily be achieved with np.rollaxis:

>>> v.dot(np.rollaxis(a, 2, 1))
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

I don't like np.dot too much, unless it is for the obvious matrix or vector multiplication, because it is very strict about the output dtype when using the optional out parameter. Joe Kington has mentioned it already, but if you are going to be doing this type of things, get used to np.einsum: once you get the hang of the syntax, it will cut down the amount of time you spend worrying about reshaping things to a minimum:

>>> a = np.ones((3, 3, 2))
>>> np.einsum('i, jki', v, a)
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

Not that it is too relevant in this case, but it is also ridiculously fast:

In [4]: %timeit a.dot(v)
100000 loops, best of 3: 2.43 us per loop

In [5]: %timeit v.dot(np.rollaxis(a, 2, 1))
100000 loops, best of 3: 4.49 us per loop

In [7]: %timeit np.tensordot(v, a, axes=(0, 2))
100000 loops, best of 3: 14.9 us per loop

In [8]: %timeit np.einsum('i, jki', v, a)
100000 loops, best of 3: 2.91 us per loop

Upvotes: 4

Joe Kington
Joe Kington

Reputation: 284562

You can also use tensordot, in this particular case. 

import numpy as np

A = np.ones((3,3,2))
v = np.array([1,2])

print np.tensordot(v, A, axes=(0, 2))

This yields:

array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

The axes=(0,2) indicates that tensordot should sum over the first axis in v and the third axis in A. (Also have a look at einsum, which is more flexible, but harder to understand if you're not used to the notation.)

If speed is a consideration, tensordot is considerably faster than using apply_along_axes for small arrays.

In [14]: A = np.ones((3,3,2))

In [15]: v = np.array([1,2])

In [16]: %timeit np.tensordot(v, A, axes=(0, 2))
10000 loops, best of 3: 21.6 us per loop

In [17]: %timeit np.apply_along_axis(v.dot, 2, A)
1000 loops, best of 3: 258 us per loop

(The difference is less apparent for large arrays due to a constant overhead, though tensordot is consistently faster.)

Upvotes: 3

NPE
NPE

Reputation: 500167

You could use numpy.apply_along_axis() for this:

In [35]: np.apply_along_axis(v.dot, 2, A)
Out[35]: 
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

The reason I think V03 doesn't work is that it's no different to:

B[:,:] = v.dot(A)

i.e. it tries to compute the dot product along the outermost axis of A.

Upvotes: 2

Related Questions