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

Question

我有一个 3D 数组，我需要在一个方向上进行转换。我有一个矩阵map_y，其中包含我使用np.dot. 是否可以避免for下面代码中的嵌套循环？import numpy as npnx, ny, nz = 64, 32, 24nyc = 11a = 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 k

map_y  :     l  x j

output : i x l      x k

因此，j对输出进行了对齐和总和减少，而其余部分是“spear-out”，为我们提供了所需的输出。

有关更多信息，请参阅链接的文档。

此外，通过将optimize标志np.einsum设置为True使用 BLAS 来处理标志。