将点积应用于 3D 数组中的所有列

Question

我有一个 3D 数组，我需要在一个方向上进行转换。我有一个矩阵map_y，其中包含我使用np.dot. 是否可以避免for下面代码中的嵌套循环？

import numpy as np

nx, ny, nz = 64, 32, 24

nyc = 11

a = np.random.rand(nz, ny, nx)

ac = np.empty((nz, nyc, nx))

map_y = np.random.rand(nyc, ny)

# Can we do this in a single numpy function?

for k in range(nz):

    for i in range(nx):

        ac[k,:,i] = np.dot(map_y, a[k,:,i])

Answer 1

使用np.einsum-ac = np.einsum('ijk,lj->ilk',a,map_y)einsum关于这个问题的一些解释：轴可以被认为是示意性的，就像这样 -a      : i      x j x kmap_y  :     l  x joutput : i x l      x k因此，j对输出进行了对齐和总和减少，而其余部分是“spear-out”，为我们提供了所需的输出。有关更多信息，请参阅链接的文档。此外，通过将optimize标志np.einsum设置为True使用 BLAS 来处理标志。