用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]]

现在，出于性能原因，我想以一种不用于循环的方式实现它。 实际代码中的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,3508,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)