python - matrix multiplication of 4 dimensional array

Question

I'm plotting a color map using a mesh grid for the map calculation. I have an X, Y gird of say 1000 by 1000 points, and some function H = function(a, b, c, X, Y) . The size of H is [2, 3, 1000, 1000] , ie for each grid point the size of H is [2, 3]. With mesh grid this is easy and efficient. Now I need to find D = np.matmul(np.transpose(H), H) . Unfortunately, I do that with 2 for loops scanning the entire grid, see code below. Can someone suggest a more elegant and efficient way to find D?

for j in range(x_mesh_length):
    for k in range(y_mesh_length):
        D[j, k] = np.matmul(H[:, :, j, k].T,H[:, :, j, k])

Answer 1

Use numpy einsum

D = np.einsum('ikml, kjml ->ijml', np.transpose(H, (1,0,2,3)), H)