将点积应用于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 k
map_y  :     l  x j
output : i x l      x k

因此，j已对齐，并且输出的总和也减少了，而其余的则是“超速输出”，为我们提供了所需的输出。

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

另外，通过将其设置为optimize以使用BLAS，在np.einsum中使用True标志。