numpy 数组比较的高效 Python 实现
Efficient Python implementation of numpy array comparisons
背景
我有两个 numpy 数组，我想用它们以尽可能 efficient/fast 的方式执行一些比较操作。两者都只包含无符号整数。
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 循环实现）这是比上面概述的方法稍慢。
如果有人有建议，我很乐意进行概要分析并报告（我已经设置了所有概要分析基础结构）。
详细说明我的评论：
展开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 的组合适用于测试用例，但可能需要针对更大的测试用例进行调整。不过idea是普遍适用的。
对于 测试的数组大小，我的解决方案相当慢。特别是它正在进行 == 测试，结果 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 的特定随机样本中，事实证明这就是一切。在 positions 中找不到任何 pairs。5000 个中有 54 个有J==1，其余为0，不匹配）。
方法 #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
背景
我有两个 numpy 数组，我想用它们以尽可能 efficient/fast 的方式执行一些比较操作。两者都只包含无符号整数。
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 循环实现）这是比上面概述的方法稍慢。
如果有人有建议，我很乐意进行概要分析并报告（我已经设置了所有概要分析基础结构）。
详细说明我的评论：
展开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 的组合适用于测试用例，但可能需要针对更大的测试用例进行调整。不过idea是普遍适用的。
对于 == 测试，结果 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 的特定随机样本中，事实证明这就是一切。在 positions 中找不到任何 pairs。5000 个中有 54 个有J==1，其余为0，不匹配）。
方法 #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