numpy数组比较的高效Python实现

时间:2015-08-08 17:40:14

标签: python performance numpy cython

背景

我有两个numpy数组,我想用它们以最有效/最快的方式进行一些比较操作。两者都只包含无符号整数。

pairs是一个n x 2 x 3数组,它包含一长串配对的3D坐标(对于某些命名法,pairs数组包含一组对......) - 即< / p>

# full pairs array
In [145]: pairs
Out[145]:
    array([[[1, 2, 4],
        [3, 4, 4]],
        .....
       [[1, 2, 5],
        [5, 6, 5]]])

# each entry contains a pair of 3D coordinates
In [149]: pairs[0]
Out[149]:
array([[1, 2, 4],
       [3, 4, 4]])

positions是一个n x 3数组,其中包含一组3D坐标

In [162]: positions
Out[162]:
array([[ 1,  2,  4],
       [ 3,  4,  5],
       [ 5,  6,  3],
       [ 3,  5,  6],
       [ 6,  7,  5],
       [12,  2,  5]])

目标 我想创建一个数组,它是pairs数组的一个子集,但是只包含条目,其中至多有一对在位置数组中 - 即不应该有两个对,其中BOTH对位于数组中。对于某些域信息,每对都将在位置列表中至少有一对位置。

到目前为止尝试的方法 我最初的天真方法是在pairs数组中循环每一对,并从positions向量中减去两对位置中的每一个,确定在两种情况下我们是否找到了由两个来自减法运算的向量中的0:

 if (~(positions-pair[0]).any(axis=1)).any() and 
    (~(positions-pair[1]).any(axis=1)).any():
    # both members of the pair were in the positions array -
    # these weren't the droids we were looking for
    pass
 else:
    # append this set of pairs to a new matrix 

这很好用,并利用一些矢量化,但有可能有更好的方法吗?

对于这个程序的其他一些对性能敏感的部分,我已经用Cython重新编写了一些内容,它带来了大量的加速,尽管在这种情况下(至少基于一个天真的嵌套for循环实现),这个稍慢比上面概述的方法。

如果有人有建议我很乐意分析和报告(我已经设置了所有的配置文件)。

2 个答案:

答案 0 :(得分:6)

方法#1

正如问题所述,两个数组都只包含无符号ints,可以利用它来将XYZ合并为一个线性索引等效版本,该版本对于每个唯一{{1}都是唯一的三重奏。实现看起来像这样 -

XYZ

由于您正在处理3D坐标,因此您还可以使用cdist检查输入数组之间的相同maxlen = np.max(pairs,axis=(0,1)) dims = np.append(maxlen[::-1][:-1].cumprod()[::-1],1) pairs1D = np.dot(pairs.reshape(-1,3),dims) positions1D = np.dot(positions,dims) mask_idx = ~(np.in1d(pairs1D,positions1D).reshape(-1,2).all(1)) out = pairs[mask_idx] 三元组。下面列出了两个实现这个想法的实现。

方法#2

XYZ

方法#3

from scipy.spatial.distance import cdist

p0 = cdist(pairs[:,0,:],positions)
p1 = cdist(pairs[:,1,:],positions)
out = pairs[((p0==0) | (p1==0)).sum(1)!=2]

运行时测试 -

mask_idx = ~((cdist(pairs.reshape(-1,3),positions)==0).any(1).reshape(-1,2).all(1))
out = pairs[mask_idx]

验证结果

In [80]: n = 5000
    ...: pairs = np.random.randint(0,100,(n,2,3))
    ...: positions= np.random.randint(0,100,(n,3))
    ...: 

In [81]: def cdist_split(pairs,positions):
    ...:    p0 = cdist(pairs[:,0,:],positions)
    ...:    p1 = cdist(pairs[:,1,:],positions)
    ...:    return pairs[((p0==0) | (p1==0)).sum(1)!=2]
    ...: 
    ...: def cdist_merged(pairs,positions):
    ...:    mask_idx = ~((cdist(pairs.reshape(-1,3),positions)==0).any(1).reshape(-1,2).all(1))
    ...:    return pairs[mask_idx]
    ...: 
    ...: def XYZ_merged(pairs,positions):
    ...:    maxlen = np.max(pairs,axis=(0,1))
    ...:    dims = np.append(maxlen[::-1][:-1].cumprod()[::-1],1)
    ...:    pairs1D = np.dot(pairs.reshape(-1,3),dims)
    ...:    positions1D = np.dot(positions,dims)
    ...:    mask_idx1 = ~(np.in1d(pairs1D,positions1D).reshape(-1,2).all(1))
    ...:    return pairs[mask_idx1]
    ...: 

In [82]: %timeit cdist_split(pairs,positions)
1 loops, best of 3: 662 ms per loop

In [83]: %timeit cdist_merged(pairs,positions)
1 loops, best of 3: 615 ms per loop

In [84]: %timeit XYZ_merged(pairs,positions)
100 loops, best of 3: 4.02 ms per loop

答案 1 :(得分:3)

阐述我的评论:

展开pairs会更有趣。随意使用更大,更逼真的阵列进行测试:

In [260]: pairs = np.array([[[1,2,4],[3,4,4]],[[1,2,5],[5,6,5]],[[3,4,5],[3,5,6]],[[6,7,5],[1,2,3]]])

In [261]: positions = np.array([[ 1,  2,  4],
       [ 3,  4,  5],
       [ 5,  6,  3],
       [ 3,  5,  6],
       [ 6,  7,  5],
       [12,  2,  5]])

将两个数组扩展为可广播的形状:

In [262]: I = pairs[None,...]==positions[:,None,None,:]

In [263]: I.shape
Out[263]: (6, 4, 2, 3)

大布尔数组,显示所有维度上的元素匹配。可以自由地替换其他比较(difference ==0np.isclose用于浮点数等)。

In [264]: J = I.all(axis=-1).any(axis=0).sum(axis=-1)

In [265]: J
Out[265]: array([1, 0, 2, 1])

整合各种维度的结果。匹配坐标上的所有数字,匹配任何位置,计算匹配数。

In [266]: pairs[J==1,...]
Out[266]: 
array([[[1, 2, 4],
        [3, 4, 4]],

       [[6, 7, 5],
        [1, 2, 3]]])

J==1表示只有一对值匹配的元素。 (见尾注)

anyandsum的组合适用于测试用例,但可能需要使用更大的测试用例进行调整。但是这个想法通常是适用的。

对于https://stackoverflow.com/a/31901675/901925测试的数组的大小,我的解决方案非常慢。特别是它正在进行==测试,结果I的形状为(5000, 5000, 2, 3)

压缩最后一个维度有很大帮助

dims = np.array([10000,100,1])  # simpler version of dims from XYZmerged
pairs1D = np.dot(pairs.reshape(-1,3),dims)
positions1D = np.dot(positions,dims)
I1d = pairs1D[:,None]==positions1D[None,:]
J1d = I1d.any(axis=-1).reshape(pairs.shape[:2]).sum(axis=-1)

我更改了J1d表达式以匹配我的 - 计算每对匹配的数量。

in1d1使用的Divakar更快:

mask = np.in1d(pairs1D, positions1D).reshape(-1,2)
Jmask = mask.sum(axis=-1)

我刚刚意识到OP要求at most one of the pairs is in the positions array。我在哪里测试exactly one match per pair因此,我的测试应全部更改为pairs[J<2,...]

(在我特定的随机样本中,n = 5000,结果证明是一切。没有任何pairspositions中都有J==1ImageView,其余为0,不匹配。