在没有重复的情况下找到两个 numpy 数组之间最近对的最快方法

Question

给定两个具有不同行数 (len(B) > len(A)) 但相同列数 (A.shape[1] = B.shape[1] = 3) 的大型 numpy 数组 A 和 B。我想知道从 B 获得子集 C 的最快方法，该子集具有最小总距离（所有成对距离的总和）到 A 没有重复项（每对必须是唯一的）。这意味着 C 应该与 A.

具有相同的形状

下面是我的代码，但有两个主要问题：

我不知道这是否给出了最小总距离
实际上我有一个比 np.linalg.norm 更昂贵的距离计算函数（需要处理周期性边界条件）。我认为这绝对不是最快的方法，因为下面的代码每次调用一对距离计算函数。当我调用更昂贵的距离计算函数时会有很大的开销，它将永远运行。有什么建议吗？

import numpy as np
from operator import itemgetter
import random
import time

A = 100.*np.random.rand(1000, 3)
B = A.copy()
for (i,j), _ in np.ndenumerate(B):
    B[i,j] += np.random.rand()
B = np.vstack([B, 100.*np.random.rand(500, 3)])

def calc_dist(x, y):
    return np.linalg.norm(x - y)

t0 = time.time()
taken = []
for rowi in A:
    res = min(((k, calc_dist(rowi, rowj)) for k, rowj in enumerate(B)
                if k not in taken), key=itemgetter(1))
    taken.append(res[0])

C = B[taken]

print(A.shape, B.shape, C.shape)
>>> (1000, 3) (1500, 3) (1000, 3)

print(time.time() - t0)
>>> 12.406389951705933

编辑：对于那些对昂贵的距离计算功能感兴趣的人，它使用ase包（可以通过pip install ase安装）

from ase.geometry import find_mic
def calc_mic_dist(x, y):
    return find_mic(np.array([x]) - np.array([y]), 
                    cell=np.array([[50., 0.0, 0.0], 
                                   [25., 45., 0.0], 
                                   [0.0, 0.0, 100.]]))[1][0]

Answer 1

利用 numpy 广播和矢量化的力量

find_mic ase.geometry 中的方法可以处理 2d np 数组。

from ase.geometry import find_mic
def calc_mic_dist(x, y):
    return find_mic(x - y, 
                    cell=np.array([[50., 0.0, 0.0], 
                                   [25., 45., 0.0], 
                                   [0.0, 0.0, 100.]]))[1]

测试：

x = np.random.randn(1,3)
y = np.random.randn(5,3)

print (calc_mic_dist(x,y).shape)
# It is a distance metrics so:
assert np.allclose(calc_mic_dist(x,y), calc_mic_dist(y,x))

输出：

(5,)

如您所见，指标是针对 x 的每个值和 y 的每个值计算的，因为 numpy 中的 x-y 具有广播的魔力。

解决方案：

def calc_mic_dist(x, y):
    return find_mic(x - y, 
                    cell=np.array([[50., 0.0, 0.0], 
                                   [25., 45., 0.0], 
                                   [0.0, 0.0, 100.]]))[1]

t0 = time.time()
A = 100.*np.random.rand(1000, 3)
B = 100.*np.random.rand(5000, 3)
selected = [np.argmin(calc_mic_dist(a, B)) for a in A]
C = B[selected]
print (A.shape, B.shape, C.shape)

print (f"Time: {time.time()-t0}")

输出：

(1000, 3) (5000, 3) (1000, 3)
Time: 9.817562341690063

在 google collab

上花费大约 10 秒

测试：

我们知道 calc_mic_dist(x,x) == 0 所以如果 A 是 B 的子集那么 C 应该正好是 A

A = 100.*np.random.rand(1000, 3)
B = np.vstack([100.*np.random.rand(500, 3), A, 100.*np.random.rand(500, 3)])
selected = [np.argmin(calc_mic_dist(a, B)) for a in A]
C = B[selected]
print (A.shape, B.shape, C.shape)
print (np.allclose(A,C))

输出：

(1000, 3) (2000, 3) (1000, 3)
True

编辑 1：避免重复

Once a vector in B is selected it cannot be again selected for other values of A

这可以通过从 B 中删除所选向量来实现，这样它就不会再次出现在 A 的下一行中作为可能的候选对象。

A = 100.*np.random.rand(1000, 3)
B = np.vstack([100.*np.random.rand(500, 3), A, 100.*np.random.rand(500, 3)])

B_ = B.copy()
C = np.zeros_like(A)

for i, a in enumerate(A):
  s = np.argmin(calc_mic_dist(a, B_))
  C[i] = B_[s]
  # Remove the paried 
  B_ = np.delete(B_, (s), axis=0)

print (A.shape, B.shape, C.shape)
print (np.allclose(A,C))

输出：

(1000, 3) (2000, 3) (1000, 3)
True

Answer 2

如果您可以计算整个 N² 距离，这对于您给出的尺寸来说并不昂贵，scipy.optimize 有一个函数可以直接解决这个问题。

import scipy.optimize
cost = np.linalg.norm(A[:, np.newaxis, :] - B, axis=2)
_, indexes = scipy.optimize.linear_sum_assignment(cost)
C = B[indexes]

在没有重复的情况下找到两个 numpy 数组之间最近对的最快方法

Fastest way to find the nearest pairs between two numpy arrays without duplicates

python

arrays

performance

numpy

distance

解决方案：

测试：

编辑 1：避免重复