r - 用于计算(平方)马氏距离的矢量化代码

Question

编辑 2:这篇文章似乎已经从 CrossValidated 转移到 StackOverflow，因为它主要是关于编程的，但这意味着花哨的 MathJax 不再起作用了。希望这仍然是可读的。

假设我想用协方差矩阵 S 计算两个向量 x 和 y 之间的平方马氏距离。这是一个由

定义的相当简单的函数

M2(x, y; S) = (x - y)^T * S^-1 * (x - y)

使用 python 的 numpy 包我可以这样做

# x, y = numpy.ndarray of shape (n,)
# s_inv = numpy.ndarray of shape (n, n)
diff = x - y
d2 = diff.T.dot(s_inv).dot(diff)

或在 R 中作为

diff <- x - y
d2 <- t(diff) %*% s_inv %*% diff

不过，就我而言，我得到了

• m by n 矩阵 X
• n维向量mu
• n by n 协方差矩阵S

并且想要找到 m 维向量 d 使得

d_i = M2(x_i, mu; S)  ( i = 1 .. m )

其中 x_i 是 X 的第 i 行。

这在 python 中使用一个简单的循环并不难:

d = numpy.zeros((m,))
for i in range(m):
    diff = x[i,:] - mu
    d[i] = diff.T.dot(s_inv).dot(diff)

当然，考虑到外循环是在 python 中而不是在 numpy 库中的 native 代码中发生的，这意味着它没有达到预期的速度。 $n$ 和 $m$ 分别约为 3-4 和几十万，我在交互式程序中经常这样做，因此加速将非常有用。

从数学上讲，我能够使用基本矩阵运算来表达这一点的唯一方法是

d = diag( X' * S^-1 * X'^T )

在哪里

 x'_i = x_i - mu

写一个矢量化版本很简单，但不幸的是，计算一个超过 100 亿的元素矩阵并且只取对角线的效率低下...我相信这个操作应该很容易用爱因斯坦符号表达，因此有望使用 numpy 的 einsum 函数进行快速评估，但我什至还没有开始弄清楚这个黑魔法是如何工作的。

所以，我想知道:有没有更好的方法来数学地表达这个操作(就简单的矩阵操作而言)，或者有人可以建议一些很好的矢量化(python 或 R)代码来有效地做到这一点？

奖金问题，勇敢者

我实际上不想做一次，我想做 k ~ 100 次。鉴于:

• m 由 n 矩阵 X

• k 由 n 矩阵U

• 由 n 个协方差矩阵组成的 n 集，每个表示为 S_j(j = 1..k)

通过 k 矩阵 D 找到 m 使得

D_i,j = M(x_i, u_j; S_j)

其中i = 1..m，j = 1..k，x_i是i X 和 u_j 的第 j 行是 U 的第 j 行。

即向量化以下代码:

# s_inv is (k x n x n) array containing "stacked" inverses
# of covariance matrices
d = numpy.zeros( (m, k) )
for j in range(k):
    for i in range(m):
        diff = x[i, :] - u[j, :]
        d[i, j] = diff.T.dot(s_inv[j, :, :]).dot(diff)

Answer 1

首先，您似乎得到了 S 然后将其取反。你不应该那样做；它很慢而且数字不准确。相反，您应该获得 S 的 Cholesky 因子 L，以便 S = L L^T；然后

M^2(x, y; L L^T)
  = (x - y)^T (L L^T)^-1 (x - y)
  = (x - y)^T L^-T L^-1 (x - y)
  = || L^-1 (x - y) ||^2,

并且由于 L 是三角形 L^-1 (x - y) 可以高效计算。

事实证明，如果您适本地 reshape scipy.linalg.solve_triangular ，它会很高兴地同时执行这些操作:

L = np.linalg.cholesky(S)
y = scipy.linalg.solve_triangular(L, (X - mu[np.newaxis]).T, lower=True)
d = np.einsum('ij,ij->j', y, y)

稍微分解一下，y[i, j] 是 L^-1 (X_j -\mu) 的第 i 个分量。然后 einsum 调用会执行

d_j = \sum_i y_{ij} y_{ij}
    = \sum_i y_{ij}^2
    = || y_j ||^2,

就像我们需要的那样。

不幸的是，solve_triangular 不会对其第一个参数进行矢量化，因此您可能应该只在此处循环。如果 k 仅为 100 左右，那不会是一个重大问题。

---

如果你实际上得到的是 S^-1 而不是 S，那么你确实可以更直接地使用 einsum 来做到这一点。由于 S 在您的情况下很小，因此实际反转矩阵然后执行此操作也可能会更快。但是，一旦 n 不是平凡的大小，这样做就会损失很多数值精度。

要弄清楚如何使用 einsum，请根据组件编写所有内容。我将直接进入奖金案例，为了符号方便，写 S_j^-1 = T_j:

D_{ij} = M^2(x_i, u_j; S_j)
  = (x_i - u_j)^T T_j (x_i - u_j)
  = \sum_k (x_i - u_j)_k ( T_j (x_i - u_j) )_k
  = \sum_k (x_i - u_j)_k \sum_l (T_j)_{k l} (x_i - u_j)_l
  = \sum_{k l} (X_{i k} - U_{j k}) (T_j)_{k l} (X_{i l} - U_{j l})

因此，如果我们制作形状为 (m, n) 的数组 X，形状为 (k, n) 的 U , 和形状为(k, n, n)的T，那么我们可以把它写成

diff = X[np.newaxis, :, :] - U[:, np.newaxis, :]
D = np.einsum('jik,jkl,jil->ij', diff, T, diff)

其中 diff[j, i, k] = X_[i, k] - U[j, k]。