Computing expectations of an array, suggestions for speed improvements

 

Questions


I have a (N0, N1, N2, N3) Matrix V and a (N1, N1) Matrix M. N1 is typically around 30-50, N0xN1xN2xN3 is around 1 000 000. I want a new Matrix EV, for which the i0, i1, i2, i3 entry is given by

np.sum(V[i0, :, i2, i3] * M[i1, :])    

My current code to achieve that is:

  V_exp = np.tile(V[:, :, :, :, None], (1, 1, 1, 1, N1))
  M_exp = np.tile(M.T[None, :, None, None, :], (N0, 1, N2, N3, 1))
  EV = np.sum(V_exp * M_exp, axis = 1)
  EV = np.rollaxis(EV, 3, 1)

I have to perform this operation many times, and it is the absolute bottleneck of my code. I wonder if there’s any potential for a speed improvement of my code. I appriciate suggestions!

 

 

————————————————-

Answer

A single call to np.einsum would perform all those operations in one go, after extending V to a 5D shape with np.newaxis/None, like so –

EV  = np.einsum('ijklm,mj->imkl',V[...,None],M)

Thus, we are avoiding any intermediate arrays for a memory efficient solution.

Explanation

(1) Starting code that involved extending dimensions :

V_exp = np.tile(V[:, :, :, :, None], (1, 1, 1, 1, N1))
M_exp = np.tile(M.T[None, :, None, None, :], (N0, 1, N2, N3, 1))
output1 = V_exp * M_exp

In np.einsum‘s way, it would be translated as :

np.einsum('ijklm,mj->ijklm',V[...,None],M)

Please notice that we have used mj instead of the usual jm for M to correspond to M.T[None, :, None, None, :].

(2) Next up, we have :

EV = np.sum(V_exp * M_exp, axis = 1)

Thus, we are summing along axis = 1, so the einsum call would need the output string specifier be changed from ->ijklm to ->iklm.

(3) Finally :

EV = np.rollaxis(EV, 3, 1)

For this porting of rollaxis, we just need to push up axis=3 to axis=1‘s position and push down each of axes=1,2 one position to the right. Thus, the output string specifier would change from ->iklm to ->imkl to give us :

np.einsum('ijklm,mj->imkl',V[...,None],M)

arrays,numpy,performance

Facebook Comments

Post a comment