Replacing for loops with numpy operations in case of 3D matrices-CodePudding

I have this mathematical formula to implement ![https://i.stack.imgur.com/S28BA.png] where for example w_fk denotes matrix of shape (F, K). I have implemented this as

gamma_dashed_lft = np.zeros((L, F, T))
for l in range(L):
    for f in range(F):
        for t in range(T):
            temp = 0
            for k in range(K):
                temp = temp   (q_lk[l, k] * w_fk[f, k] * h_kt[k, t])
            gamma_dashed_lft[l, f, t] = temp
return gamma_dashed_lft

What would be the way to replace for loops with matrix multiplication in the case of given formula?

CodePudding user response：

The product q[l,k]*w[f,k] is evaluated for every value of t, so I would approach this as:

for l,k:
  qw[k] = q[l,k] * w[f,k]

and then use that temporary array qw in the inner loop. Now you have a bunch of inner products, in effect a matrix-vector product, and you've saved a factor of T in operations.

However, that does not give you the cache optimization that a matrix-matrix multiplication gives. For that, create a matrix qw_l for each l:

for l:
  qw_l[f,k] = q[l,k]*w[f,k]
  matrix-matrix-product: qw_l times h

(Note that I'm not spelling out all the loops!)

This costs you one array of size F x K extra but that's probably not a problem. In case you're wondering, the cost of creating it is F x K, while the matrix-matrix product is F x T x K, so you don't have to worry about that.

CodePudding user response：

You should have provided a concrete example. Fortunately it's not too hard to read the dimensions and create:

In [302]: L,F,T,K=2,3,4,5
In [303]: q_lk=np.arange(L*K).reshape(L,K)
In [304]: w_fk=np.arange(F*K).reshape(F,K)
In [305]: h_kt=np.arange(K*T).reshape(K,T)

which when applied to your code produces:

In [306]: gamma_dashed_lft = np.zeros((L, F, T))
     ...: for l in range(L):
     ...:     for f in range(F):
     ...:         for t in range(T):
     ...:             temp = 0
     ...:             for k in range(K):
     ...:                 temp = temp   (q_lk[l, k] * w_fk[f, k] * h_kt[k, t])
     ...:             gamma_dashed_lft[l, f, t] = temp
     ...: 

In [308]: gamma_dashed_lft
Out[308]: 
array([[[ 400.,  430.,  460.,  490.],
        [1000., 1080., 1160., 1240.],
        [1600., 1730., 1860., 1990.]],

       [[1000., 1080., 1160., 1240.],
        [2600., 2855., 3110., 3365.],
        [4200., 4630., 5060., 5490.]]])

An equivalent expression making full use of broadcasting is:

In [309]: arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,:]).sum(axis=2)
In [310]: arr.shape
Out[310]: (2, 3, 4)
In [311]: np.allclose(arr,gamma_dashed_lft)
Out[311]: True

In setting up the broadcasting I was aiming for an array with shape (L,F,K,T) with the sum reduction on the K.

Since you made me create the test case, I let you work out the broadcasting details. It'll be a good exercise for you.