python - 缩放/旋转成对平方欧氏距离的矢量化计算

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)

此外，还可以使用 optimize标记于np.einsum将其设置为 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变成np.einsum('ijk,kl,ijl->ij',V,A,V)使用einsum的字符串表示法。更多信息可以在 docs 找到.