python - NumPy - 沿第三维点积而不复制

Question

我正在尝试向量化一个函数，该函数将 3 分量向量“x”和 3x3“矩阵”作为输入并生成标量

def myfunc(x, matrix):
    return np.dot(x, np.dot(matrix, x))

但是这需要被调用“n”次，并且向量 x 每次都有不同的分量。我想修改这个函数，使其接受一些 3xn 数组(其中的列是向量 x)作为输入，并生成一个向量，其分量是在每次迭代时计算的标量。

我可以写出一个爱因斯坦求和来完成这项工作，但它要求我构造一个 3x3xn 的原始 3x3 的“副本”堆栈。我担心这样做会抵消我尝试这样做所获得的任何性能提升。有没有办法计算我想要的向量而不复制 3x3？

Answer 1

令 x 为 3xN 数组，y 为 3x3 数组。您正在寻找

z = numpy.einsum('ji,jk,ki->i', x, y, x)

您还可以将您正在讨论的 3x3xN 数组构建为 y 的 View 以避免复制，但这不是必需的。