numpy数组比较的高效Python实现

Question

背景

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

pairs 是一个 n x 2 x 3 数组，其中包含一长串成对的 3D 坐标（对于某些命名法，pairs 数组包含一组对...） - 即

# 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 数组的子集，但仅包含最多一对在位置数组中的条目 - 即不应该有两个对都位于位置数组中的对。对于某些域信息，每对将在位置列表中至少有一个对位置。

目前尝试的方法 我最初的幼稚方法是循环遍历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 

这很好用，并利用了 some 矢量化，但可能有更好的方法来做到这一点？

对于这个程序的其他一些对性能敏感的部分，我用 Cython 重写了一些东西，这带来了巨大的加速，但在这种情况下（至少基于一个简单的嵌套 for 循环实现），这会稍微慢一些比上面概述的方法。

如果人们有任何建议，我很乐意进行分析并反馈（我已设置所有分析基础架构）。

Answer 1

详细说明我的评论：

扩展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 ==0、np.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 表示该对中只有一个值匹配的元素。 （见尾注）

any、and 和 sum 的组合适用于测试用例，但可能需要针对更大的测试用例进行调整。但想法是普遍适用的。

---

对于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 表达式以匹配我的 - 计算每对匹配的数量。

Divakar 使用的in1d1 更快：

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 的特定随机样本中，结果证明是一切。没有任何 pairs 两者都在 positions 中找到。5000 个中有 54 个有 J==1，其余的为 0，不匹配）。

Answer 2

方法#1

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

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]

由于您正在处理 3D 坐标，您还可以使用cdist 来检查输入数组之间的相同XYZ 三元组。接下来列出的是两个考虑到这个想法的实现。

方法 #2

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]

方法#3

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

验证结果 -

In [85]: np.allclose(cdist_split(pairs,positions),cdist_merged(pairs,positions))
Out[85]: True

In [86]: np.allclose(cdist_split(pairs,positions),XYZ_merged(pairs,positions))
Out[86]: True