Python 将点积应用于三维阵列中的所有列

我有一个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])

使用-

关于问题的

einsum

说明：

轴可以被认为是示意图，就像这样-

a      : i      x j x k
map_y  :     l  x j
output : i x l      x k

因此，

j

是对齐的，并且输出的和也减少了，而其余的是“叉出来的”，为我们提供了所需的输出

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

---

另外，在

np.einsum

中将

optimize

标记设置为

True

以使用BLAS，可以使用

optimize

来使用它。

您能更详细地解释一下

einsum

的第一个参数吗？如果将

nz，ny，nx，nyc

映射到

i，j，k，l

，则

中的
字符串几乎会自动写入。它们是
a
和
map_y
（和
ac
）各自维度的索引。Chiel补充了一些解释。发布的解决方案对你有用吗？效果很好。我不知道
einsum
的功能。
a      : i      x j x k
map_y  :     l  x j
output : i x l      x k