缩放/旋转成对平方欧几里德距离的矢量化计算

Question

给定一组 n 个维度为 d 的向量存储在 (n,d) 数组中，以及第二组 m 个相同维度的向量（存储在 (m,d) 数组中）我想计算向量之间的平方点距离，由一些矩阵 A 缩放，大小为 (d,d).

输出应该是一个(n,m)数组。

我希望 m 和 n 的输入范围在 1 到 10.000 之间，d 的输入范围在 1 到 100 之间。

两点之间的距离由下式给出：

在未优化但工作的 Python 代码中，如下所示：

import numpy as np

v1 = np.array([[1, 2],
               [3, 4],
               [4, 5]])

v2 = np.array([[1,1],
               [2, 2],
               [2, 2],
               [0, 0]])

A = np.array([[1,0], [2, 3]])

d = np.zeros((3, 4))

for i in range(0,3):
    for j in range(0,4):
        d[i,j] = (v1[i,:] - v2[j,:]).T @ A @ (v1[i,:] - v2[j,:])

示例点之间的平方距离为：

d = [[  3.   1.   1.  17.]
 [ 43.  17.  17.  81.]
 [ 81.  43.  43. 131.]]

是否有这样的版本，可以避免 python 中的嵌套循环，例如使用广播黑魔法？

编辑：

案例

A = np.array([[1,0], [0, 1]])

这是可以计算的正常平方欧几里得距离，例如

from scipy.spatial.distance import cdist

cdist(v1,v2,'sqeuclidean')

Answer 1

我们可以使用np.einsum -

V = v1[:,None,:]-v2
d_out = np.einsum('ijk,kl,ijl->ij',V,A,V)

另外，通过将np.einsum 中的optimize 标志设置为True 来使用BLAS。

矢量化方法说明

原始代码是-

d[i,j] = (v1[i,:] - v2[j,:]).T @ A @ (v1[i,:] - v2[j,:])

我。我们正在翻译：

v1[i,:] - v2[j,:]

到外部操作用broadcasting：

v1[:,None,:]-v2

示意图：

v1[:,None,:]  :  m x 1 x n
v2            :      m x n
output, V     :  m x m x n

More info on outer explanation.

更多关于broadcasting的信息可以在docs找到。

二。接下来，(v1[i,:] - v2[j,:]).T @ A @ (v1[i,:] - v2[j,:]) 和新的 V 使用 einsum 的字符串表示法变为 np.einsum('ijk,kl,ijl->ij',V,A,V)。更多信息可以在docs找到。