Vectorization of a large number of NumPy

Question

I am interested in calculating a large NumPy array. I have a large array A that contains a bunch of numbers. I want to calculate the sum of various combinations of these numbers. The data structure is as follows:

 A = np.random.uniform(0,1, (3743, 1388, 3)) Combinations = np.random.randint(0,3, (306,3)) Final_Product = np.array([ np.sum( A*cb, axis=2) for cb in Combinations]) 

My question is, is there a more elegant and efficient way to count data? I find it difficult to work with np.dot() when a three-dimensional array is involved.

If this helps, the Final_Product form Final_Product ideally be (3743, 306, 1388). Final_Product has the form (306, 3743, 1388), so I can just redo it to get there.

Answer 1

np.dot() will not give you the desired result unless you include additional steps, which probably include reshaping is one vectorized approach using np.einsum to do this in one shot without unnecessary memory overhead -

 Final_Product = np.einsum('ijk,lk->lij',A,Combinations) 

For completeness, here is np.dot and reshaping , as discussed earlier -

 M,N,R = A.shape Final_Product = A.reshape(-1,R).dot(Combinations.T).T.reshape(-1,M,N) 

Run-time checks and output checks -

 In [138]: # Inputs ( smaller version of those listed in question ) ...: A = np.random.uniform(0,1, (374, 138, 3)) ...: Combinations = np.random.randint(0,3, (30,3)) ...: In [139]: %timeit np.array([ np.sum( A*cb, axis=2) for cb in Combinations]) 1 loops, best of 3: 324 ms per loop In [140]: %timeit np.einsum('ijk,lk->lij',A,Combinations) 10 loops, best of 3: 32 ms per loop In [141]: M,N,R = A.shape In [142]: %timeit A.reshape(-1,R).dot(Combinations.T).T.reshape(-1,M,N) 100 loops, best of 3: 15.6 ms per loop In [143]: Final_Product =np.array([np.sum( A*cb, axis=2) for cb in Combinations]) ...: Final_Product2 = np.einsum('ijk,lk->lij',A,Combinations) ...: M,N,R = A.shape ...: Final_Product3 = A.reshape(-1,R).dot(Combinations.T).T.reshape(-1,M,N) ...: In [144]: print np.allclose(Final_Product,Final_Product2) True In [145]: print np.allclose(Final_Product,Final_Product3) True 

Answer 2

Instead of dot you can use tensordot . Your current method is equivalent to:

 np.tensordot(A, Combinations, [2, 1]).transpose(2, 0, 1) 

Note the transpose at the end to place the axes in the correct order.

Like dot , the tensordot function can access the fast BLAS / LAPACK libraries (if you installed them) and therefore should work well for large arrays.