python-4维数组的矩阵乘法

Question

我正在使用网格图来绘制颜色图以进行图计算。我有一个X，Y围栏，例如1000 x 1000点，以及一些函数H = function(a, b, c, X, Y)。 H的大小为[2, 3, 1000, 1000]，即，对于每个网格点，H的大小为[2，3]。使用网格，这既简单又有效。 现在，我需要找到D = np.matmul(np.transpose(H), H)。不幸的是，我使用2 for循环来扫描整个网格，请参见下面的代码。有人可以建议一种更优雅，更有效的方法来找到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

使用numpy einsum

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