基准

Question

我有一个长度为embed_vec的numpy数组tot_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

我对这个问题很感兴趣，并尝试使用scipy的https://regex101.com/r/YIJh9g/4来有效地解决它。然而，这种方法可能耗尽内存，因为在内部保持距离＆lt; = R的所有对的列表。如果您的R和tot_vec足够小，它将起作用：

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性能。

Answer 3

首先播出差异：

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是它们去向的布尔掩码。你可以从那里进行处理。

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

3 个答案:

基准