我有一个名为 vec 的矩阵，有两列:vec[:,0] 和 vec[:,1]。 P 包含两个矩阵:P[0,:,:] 和 P[1,:,:]。我想将 P[0,:,:] 与 vec 的第一列相乘，并将 P[1,:,:] 与 vec 的第二列相乘。然而，操作 P@vec 还给出了 P[0,:,:] 与 vec 第二列的矩阵乘积以及 P[1,:,:] 与 vec 第一列的矩阵乘积，这会减慢速度我的代码。
是否可以直接计算第 1 列到矩阵 1 以及第 2 列到矩阵 2 的对，而无需“关闭”乘积？
import numpy as np
P=np.arange(50).reshape(2, 5, 5)
vec=np.arange(10).reshape(5,2)
have=P@vec
want=np.column_stack((have[0,:,0],have[1,:,1]))
have,want
最佳答案
numpy 中有一个非常强大的函数，称为np.einsum。它可以执行各种张量收缩、轴重新排序和矩阵乘法。对于您的示例，您可以使用
res = np.einsum('nij,jn->in', P, vec)
之后的 res 与 want 完全相同。
这是如何工作的:
您为 np.einsum 函数提供数组和签名('nij,jn->in' 字符串)，告诉函数如何将数组相乘。简而言之，您希望 P 张量的第三个轴与 vec 第一个轴收缩。因此，您在签名字符串中选择相同的索引j，并将其保留在->之后的部分中。如果索引出现在 -> 的左侧和右侧，则仅完成广播，此处针对 n 和 i 进行广播指数。
关于这个非常强大的功能的更完整的解释以及如何使用它的许多示例可以在 the corresponding numpy documentation 找到。 .
关于python - 3d 数组到矩阵乘法，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/70655294/