沿数组轴进行矩阵向量乘法

Question

在一个当前项目中，我有一个形状为（I，J，K，N）的大型多维数组和一个维度为N的方阵。

我需要对数组的最后一个轴执行与方阵的矩阵向量乘法。

因此，显然的解决方案是：

for i in range(I):
    for j in range(J):
        for k in range(K):
            arr[i,j,k] = mat.dot(arr[i,j,k])

但是这种方法速度相对较慢。因此，我也尝试了NumPy的tensordot函数，但成果有限。 我期望像这样：

arr = tensordot(mat,arr,axes=((0,1),(3))) 

应该可以解决问题，但我收到了形状不匹配的错误。

有没有更好的解决方法或者知道如何正确使用tensordot？

谢谢！

- cp3028

Answer 1

这应该会完成你的循环，但使用向量化循环：

from numpy.core.umath_tests import matrix_multiply

arr[..., np.newaxis] = matrix_multiply(mat, arr[..., np.newaxis])

matrix_multiply 和它的姐妹函数 inner1d 是numpy中隐藏的、未记录文档的宝石，虽然numpy 1.8版本将会有一整套的线性代数gufuncs。 matrix_multiply 对于输入数组的最后两个维度进行矩阵乘法，并在其余维度上进行广播。唯一棘手的部分是设置一个额外的维度，以便在乘法时看到列向量，并且在重新分配回数组时也添加这个维度，以避免形状不匹配。

- Jaime

---

谢谢，这正是我需要在多个向量上进行矩阵乘法的方法。 我使用了 xyzRot = matrix_multiply(rot,xyz[...,np.newaxis]).reshape(xyz.shape)，其中 rot 是一个 (3,3) 的旋转矩阵，而 xyz 是一个形状为 (n,3) 的向量数组。非常有用！ - letmaik

Answer 2

我认为你的for循环有误，对于这种情况，dot似乎已经足够了。

# a is your IJKN
# b is your NN
c = dot(a, b)

这里的c将是一个IJKN数组。如果您想要对最后一个维度求和以获取IJK数组：

arr = dot(a,b).sum(axis=3)

但我不确定这是否是您想要的…

- Saullo G. P. Castro