python - 乘法矩阵A乘以多维矩阵 "matrix-wise?"

Question

我有一个 3d numpy 数组 u，形状 (k, m, n)，我正在尝试计算一个新数组 uprod，形状 (k, m, n)，这样 uprod[j] = np.dot (A, u[j])，其中 A 是一个完全不依赖于 j 的固定矩阵。我可以很容易地在最里面的索引上写一个循环，但是有更快/更好的方法吗？

Answer 1

np.einsum可以胜任这项工作:

result = numpy.einsum('ij,kjl->kil', A, u)

您还可以使其广播，因此如果 X 和 Y 被认为是二维矩阵数组，则以下调用将执行适当广播的 点:

result = numpy.einsum('...ij,...jk->...ik', X, Y)

例如，如果 X 的形状为 (3, 4, 5, 6) 而 Y 的形状为 (4, 6, 5)，那么 result[1, 2] 将是形状为 (5, 5) 的数组，等于 X[1, 2].dot(Y[2]).

---

您也可以使用 dot 来做到这一点. A.dot(u) 生成一个结果数组，其中 A.dot(u)[i, j, k] == A[i, :].dot(u[j, : , k)。你想要一个结果数组 result[i, j, k] == A[j, :].dot(u[i, :, k]);你可以用 rollaxis 得到这个或 transpose使用 axes 参数。

result = numpy.rollaxis(A.dot(u), 1)

在 einsum 使广播变得容易的地方，用于高维输入的 dot 有点像外积。使用与之前相同的 X 和 Y，如果您设置

result = numpy.rollaxis(X.dot(Y), axis=X.ndim-2, start=X.ndim+Y.ndim-3)

然后 result[1, 2, 3] 将是形状为 (5, 5) 的数组，等于 X[1, 2].dot (Y[3]).