用 numpy 批量计算点积

Question

我正在尝试计算坐标变换的协方差矩阵（从球面坐标到笛卡尔坐标）。该数据集包含 504 个不同的球坐标。

我使用我使用矩阵乘积实现的误差传播来计算变换的协方差矩阵。

这个矩阵产品是A dot CR dot B。这些矩阵的一般形状可以在这里看到：

A = [[dx/dr, dx/dt, dx/dp]
     [dy/dr, dy/dt, dy/dp]
     [dz/dr, dz/dt, dz/dp]]

B = [[dx/dr, dy/dr, dz/dr]
     [dx/dt, dy/dt, dz/dt]
     [dx/dp, dy/dp, dz/dp]]

CR = [[0.001 ** 2, 0, 0]
      [0, 0.003 ** 2, 0]
      [0, 0, 0.005 ** 2]]

现在我想以一种出于性能原因不使用 for 循环的方式来实现它。 实际代码中的 A 和 B（在帖子的底部）是包含 504、3x3 矩阵的数组。 因此阵列的形状是 (508, 3, 3)。我想要做的是对所有相应的矩阵A[0] dot CR dot B[0], A[0] dot CR dot B[0], ... 采用上述点积，并让结果数组的形状为 (508, 3, 3)。我目前已经使用列表推导实现了这一点（请参阅 CX 的计算），但我想完全删除 for 循环。

我尝试将 (508, 3, 3) 数组传递给 np.dot() -> CX = np.dot(np.dot(A, CR), B) 但这会导致 (508, 3, 508, 3) 数组，这不是我想要的。似乎没有关键字参数来指定应该将哪些轴用于点积，就像np.transpose() 一样。有谁知道我如何实现这个？

SIGMA_R, SIGMA_THETA, SIGMA_PHI = 0.001, 0.003, 0.005
CR = np.array([[SIGMA_R ** 2, 0, 0],
               [0, SIGMA_THETA ** 2, 0],
               [0, 0, SIGMA_PHI ** 2]])

A = np.transpose(np.array(([derivative_r(r, theta, phi),
                            derivative_theta(r, theta, phi),
                            derivative_phi(r, theta, phi)])))

B = np.transpose(A, axes=(0, 2, 1))
CX = np.array([np.dot(np.dot(A[i], CR), B[i]) for i in range(504)])

Answer 1

看来你需要np.einsum

CX = np.einsum('ijk, kl, ilm -> ijm', A, CR, B) 

甚至：

CX = np.einsum('ijk, kl, iml -> ijm', A, CR, A)