[Python] How to do this calculation faster?

import numpy as np

N=1000
M=500
L=4
O=10

C=np.random.randn(M)
IDX=np.random.randint(L,size=(N,O))
U=np.random.randn(M,N,L,L)

H=np.zeros((M,N,N))
for n1 in range(N):
idx1=IDX[n1]
for n2 in range(N):
idx2=IDX[n2]
H[:,n1,n2]=np.prod(U[:,range(O),idx1,idx2],axis=1)


I want to make calculation of H more faster and tried ray library to do this as multiprocessing, but time to generate remote list took longer time. How do you think to optimize this code?

Continue reading...