Efficient product of 1D array and 3D array along one dimension - NumPy

Question

I have two numpy arrays:

• A 1D array called t of shape (70L,) with element called let s say ti
• A 3D array called I with shape (70L, 1024L, 1024L), with each elements called Ii. Ii are thus of dimension (1024L, 1024L)

I would like to make a product of the two array along the first dimension, i.e.:

tI = t1*I1,t2*I2,...,tN*IN

such as to obtain again a new array of dimension (70L, 1024L, 1024L) and then take the sum along the first dimension in order to obtain an array of dimension (1024L, 1024L):

tsum = t1*I1 + t2*I2 + ... +tN*IN

For the moment I am satisfied with doing the following:

tI = np.asarray([t[i]*I[i,:,:] for i in range(t.shape[0])])
tsum = np.sum(tI,axis=0)

But it is going to be a bit slow is the dimensions of my array are increasing. I was wondering if there exist a numpy or scipy function, more optimized for that particular task?

Thanks in advance of any link or information.

Greg

Answer 1

You can use np.tensordot -

np.tensordot(t,I, axes=([0],[0]))

You can also use np.einsum -

np.einsum('i,ijk->jk',t,I)

Runtime test and output verification -

In [21]: def original_app(t,I):
    ...:     tI = np.asarray([t[i]*I[i,:,:] for i in range(t.shape[0])])
    ...:     tsum = np.sum(tI,axis=0)
    ...:     return tsum
    ...: 

In [22]: # Inputs with random elements
    ...: t = np.random.rand(70,)
    ...: I = np.random.rand(70,1024,1024)
    ...: 

In [23]: np.allclose(original_app(t,I),np.tensordot(t,I, axes=([0],[0])))
Out[23]: True

In [24]: np.allclose(original_app(t,I),np.einsum('i,ijk->jk',t,I))
Out[24]: True

In [25]: %timeit np.tensordot(t,I, axes=([0],[0]))
1 loops, best of 3: 110 ms per loop

In [26]: %timeit np.einsum('i,ijk->jk',t,I)
1 loops, best of 3: 201 ms per loop

Answer 2

Divakar gives the best (most efficient) answers. For completeness' sake, one other way of doing it is by using Numpy's broadcasting capabilities:

(t[:,np.newaxis,np.newaxis]*I).sum(axis=0)

By adding two axes to t, broadcasting becomes possible and one can use regular Numpy operations, which for some might be more readable.