如何有效地比较numpy数组中的条目？

Question

我有一个长度为 tot_vec 的 numpy 数组 embed_vec，其中每个条目都是一个 3d 向量：

[[ 0.52483319  0.78015841  0.71117216]
 [ 0.53041481  0.79462171  0.67234534]
 [ 0.53645428  0.80896727  0.63119403]
 ..., 
 [ 0.72283509  0.40070804  0.15220522]
 [ 0.71277758  0.38498613  0.16141834]
 [ 0.70221445  0.36918032  0.17370776]]

对于这个数组中的每个元素，我想找出与该条目“接近”的其他条目的数量。接近，我的意思是两个向量之间的距离小于指定值R。为此，我必须将这个数组中所有可能的对相互比较，然后找出数组中每个向量的接近向量的数量。所以我这样做：

p = np.zeros(tot_vec) # This contains the number of close vectors
for i in range(tot_vec-1):
    for j in range(i+1, tot_vec):
        if np.linalg.norm(embed_vec[i]-embed_vec[j]) < R:
            p[i] += 1

但是，这是非常低效的，因为我有两个嵌套的 python 循环并且对于更大的数组大小，这需要很长时间。如果这是在 C++ 或 Fortran 中，那将不是一个大问题。我的问题是，可以使用某种矢量化方法有效地使用 numpy 实现相同的目标吗？作为旁注，我也不介意使用 Pandas 的解决方案。

Answer 1

方法#1：矢量化方法 -

def vectorized_app(embed_vec, R):  
    tot_vec = embed_vec.shape[0]          
    r,c = np.triu_indices(tot_vec,1)
    subs = embed_vec[r] - embed_vec[c]
    dists = np.einsum('ij,ij->i',subs,subs)
    return np.bincount(r,dists<R**2,minlength=tot_vec)

方法#2：循环复杂度较低（对于非常大的数组）-

def loopy_less_app(embed_vec, R):  
    tot_vec = embed_vec.shape[0]
    Rsq = R**2
    out = np.zeros(tot_vec,dtype=int)
    for i in range(tot_vec):
        subs = embed_vec[i] - embed_vec[i+1:tot_vec]
        dists = np.einsum('ij,ij->i',subs,subs)
        out[i] = np.count_nonzero(dists < Rsq)
    return out

---

基准测试

原始方法-

def loopy_app(embed_vec, R):
    tot_vec = embed_vec.shape[0]
    p = np.zeros(tot_vec) # This contains the number of close vectors
    for i in range(tot_vec-1):
        for j in range(i+1, tot_vec):
            if np.linalg.norm(embed_vec[i]-embed_vec[j]) < R:
                p[i] += 1
    return p                

时间安排 -

In [76]: # Sample random array
    ...: embed_vec = np.random.rand(3000,3)
    ...: R = 0.5
    ...: 

In [77]: %timeit loopy_app(embed_vec, R)
1 loops, best of 3: 50.5 s per loop

In [78]: %timeit loopy_less_app(embed_vec, R)
10 loops, best of 3: 143 ms per loop

350x+ 加速！

建议使用更大的数组 loopy_less_app -

In [81]: # Sample random array
    ...: embed_vec = np.random.rand(20000,3)
    ...: R = 0.5
    ...: 

In [82]: %timeit loopy_less_app(embed_vec, R)
1 loops, best of 3: 4.47 s per loop

Answer 2

先广播区别：

disp_vecs=tot_vec[:,None,:]-tot_vec[None,:,:]

现在，根据您的数据集有多大，您可能希望在不进行所有数学运算的情况下进行拳头传递。如果距离小于r，则所有组件都应小于r

first_mask=np.max(disp_vec, axis=-1)<r

然后进行实际计算

disps=np.linlg.norm(disp_vec[first_mask],axis=-1)
second_mask=disps<r

现在重新分配

disps=disps[second_mask]
first_mask[first_mask]=second_mask

disps 现在是好的值，first_mask 是它们去向的布尔掩码。您可以从那里处理。

Answer 3

我对这个问题很感兴趣，并尝试使用 scipy 的cKDTree 有效地解决它。但是，这种方法可能会耗尽内存，因为在内部维护了距离

import numpy as np
from scipy.spatial import cKDTree as KDTree

tot_vec = 60000
embed_vec = np.random.randn(tot_vec, 3)
R = 0.1

tree = KDTree(embed_vec, leafsize=100)
p = np.zeros(tot_vec)
for pair in tree.query_pairs(R):
    p[pair[0]] += 1
    p[pair[1]] += 1

如果内存是个问题，可以通过一些努力将query_pairs 重写为 Python 中的生成器函数，但会牺牲 C 的性能。