从 Numpy 中的 N 个向量中找到所有唯一的(几乎)平行 3d 向量对的最快方法

Posted

技术标签:

【中文标题】从 Numpy 中的 N 个向量中找到所有唯一的(几乎)平行 3d 向量对的最快方法【英文标题】:Fastest way to find all unique pairs of (nearly) parallel 3d vectors from N vectors in Numpy 【发布时间】:2022-01-04 11:28:16 【问题描述】:

我有一个N = 10000 3d 向量的大矩阵。为简化起见,我这里以一个 10 x 3 的矩阵为例:

import numpy as np
A = np.array([[1.2, 2.3, 0.8],
              [3.2, 2.1, 0.5],
              [0.8, 4.4, 4.4],
              [-0.2, -1.1, -1.1],
              [2.4, 4.6, 1.6],
              [0.5, 0.96, 0.33],
              [1.1, 2.2, 3.3],
              [-2.2, -4.41, -6.62],
              [3.4, 5.5, 3.8],
              [-5.1, -28., -28.1]])

我想找到几乎彼此平行的所有唯一向量对。需要使用公差测量,我想获得所有唯一的行索引对(无论顺序如何)。我设法编写了以下代码:

def all_parallel_pairs(A, tol=0.1):
    res = set()
    for i, v1 in enumerate(A):
        for j, v2 in enumerate(A):
            if i == j:
                continue
            norm = np.linalg.norm(np.cross(v1, v2))
            if np.isclose(norm, 0., rtol=0, atol=tol):
                res.add(tuple(sorted([i, j])))
    return np.array(list(res))
print(all_parallel_pairs(A, tol=0.1))

out[1]: [[0 4]
         [2 3]
         [6 7]
         [4 5]
         [0 5]]

但是,由于我使用了两个 for 循环,所以当 N 很大时它会变慢。我觉得应该有更有效和 Numpyic 的方法来做到这一点。有什么建议吗?

【问题讨论】:

【参考方案1】:

请注意,函数 np.cross 接收来自文档的向量数组:

返回两个(数组)向量的叉积。

所以一种方法是使用 numpy 高级索引来找到需要计算叉积的正确向量:

# generate the i, j indices (note that only the upper triangular matrices of indices is needed)
rows, cols = np.triu_indices(A.shape[0], 1)

# find the cross products using numpy indexing on A, and the np.cross can take array of vectors
cross = np.cross(A[rows], A[cols])

# find the values that are close to 0
arg = np.argwhere(np.isclose(0, (cross * cross).sum(axis=1) ** 0.5, rtol=0, atol=0.1))

# get the i, j indices where is 0
res = np.hstack([rows[arg], cols[arg]])

print(res)

输出

[[0 4]
 [0 5]
 [2 3]
 [4 5]
 [6 7]]

表达式:

(cross * cross).sum(axis=1) ** 0.5

是在向量数组上应用np.linalg.norm 的更快替换。

【讨论】:

【参考方案2】:

作为对Dani Masejo answer 的改进更新,您可以使用 GPU_aided 或 TPU_aided 库,例如 JAX

from jax import jit

@jit
def test_jit():
    rows, cols = np.triu_indices(A.shape[0], 1)
    cross = np.cross(A[rows], A[cols])
    arg = np.argwhere(np.isclose(0, (cross * cross).sum(axis=1) ** 0.5, rtol=0, atol=0.1))
    res = np.hstack([rows[arg], cols[arg]])

    return res

print(test_jit())

使用 google colab TPU 运行时的结果如下:

100 loops, best of 5: 12.2 ms per loop       # the question code
100 loops, best of 5: 152 µs per loop        # Dani Masejo code
100 loops, best of 5: 81.5 µs per loop       # using jax library

当数据量增加时差异会很大。

【讨论】:

以上是关于从 Numpy 中的 N 个向量中找到所有唯一的(几乎)平行 3d 向量对的最快方法的主要内容,如果未能解决你的问题,请参考以下文章

如何从一个1d Numpy数组的所有排列组合中删除所有的圆台排列组合?

Python Numpy中的几个矩阵乘法

(n,) 在 numpy 和向量的上下文中是啥意思?

第四十篇 Numpy.array的基本操作——向量及矩阵的运算

从具有重复元素的向量生成所有唯一组合

找到唯一二进制向量数量的有效解决方案