Calculating sum without comprehension [duplicate]

Question

Consider the two toy arrays below:

import numpy as np
k = np.random.randint(1, 25, (5, 2, 3))
l = np.random.randint(25, 50, (7, 3))

In [27]: k
Out[27]: 
array([[[14, 15, 24],
        [21, 24,  5]],

       [[22, 19,  9],
        [21,  1, 11]],

       [[ 1, 23,  5],
        [16, 14,  2]],

       [[ 7,  3, 16],
        [23,  2,  8]],

       [[12, 24,  4],
        [ 2, 15, 20]]])

In [28]: l
Out[28]: 
array([[47, 31, 42],
       [28, 27, 26],
       [45, 32, 49],
       [29, 34, 32],
       [40, 36, 25],
       [44, 27, 31],
       [27, 35, 26]])

I can get the multiplicative sum that I am interested in as follows:

f = np.array([np.sum( k * x, axis = 2) for x in l])

In [29]: f
Out[29]: 
array([[[2131, 1941],
        [2001, 1480],
        [ 970, 1270],
        [1094, 1479],
        [1476, 1399]],

       [[1421, 1366],
        [1363,  901],
        [ 779,  878],
        [ 693,  906],
        [1088,  981]],

       [[2286, 1958],
        [2039, 1516],
        [1026, 1266],
        [1195, 1491],
        [1504, 1550]],

       [[1684, 1585],
        [1572,  995],
        [ 971, 1004],
        [ 817,  991],
        [1292, 1208]],

       [[1700, 1829],
        [1789, 1151],
        [ 993, 1194],
        [ 788, 1192],
        [1444, 1120]],

       [[1765, 1727],
        [1760, 1292],
        [ 820, 1144],
        [ 885, 1314],
        [1300, 1113]],

       [[1527, 1537],
        [1493,  888],
        [ 962,  974],
        [ 710,  899],
        [1268, 1099]]])

How can I calculate this sum without resorting to comprehension?

Answer 1

This is a good use case for np.einsum:

np.einsum('ijk,lk->lij', k, l)

---

list_comp = np.array([np.sum( k * x, axis = 2) for x in l])
np.allclose(np.einsum('ijk,lk->lij', k, l), list_comp)
# True

---

Or using broadcasting:

(l[:,None,None]*k).sum(-1)

---

Although from a quick check on timings np.einsum runs about 3 times faster