python - 将点积应用于 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 是对齐的，并且还对输出进行了求和，而其余部分则“speard-out”，为我们提供了所需的输出。

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

此外，还可以通过将 np.einsum 中的 optimize 标志设置为 True 来使用 BLAS。