高效计算 numpy 数组中行三元组上的三个点之间的角度

Question

假设我们有一个大小为

A

的 numpy 数组

M x N

，我们将其解释为维度为

M

的

N

向量。对于三个向量

a,b,c

，我们想计算它们形成的角度的余弦：

cos(angle(a,b,c)) = np.dot((a-b)/norm(a-b), (c-b)/norm(c-b))

我们想要计算

A

三元组的数量，其中应该有

(M choose 2)*(M-2)

unique 三元组（通过

a

和

c

的对称性；如果我错了，请纠正我）。我们当然可以通过三重嵌套 for 循环来完成此任务，但我希望这可以以矢量化的方式完成。我很确定我可以使用一些广播技巧来计算一个数组，该数组包含所需的输出等，但我希望有人可以提供一个方法来计算精确唯一的数量，最好不需要额外的计算。谢谢。

编辑。为了完整性，天真的暗示。带循环：

angles = []
for i in range(len(A)):
    for j in range(len(A)):
        for k in range(i+1, len(A)):
            if j not in (i,k):
                d1 = A[i] - A[j]
                d2 = A[k] - A[j]
                ang = np.dot(d1/np.linalg.norm(d1), d2/np.linalg.norm(d2))
                angles.append(ang)

Answer 1

其中应该有

(M choose 2)*(M-2)

独特的三元组（通过a和c的对称性；如果我错了，请纠正我）

我认为是对的。我数了一下

M * ((M-1) choose 2)

，是等价的。

我希望有人能够提供一个精确计算独特数量的配方，最好不需要额外的计算。

好吧，让我们从简单的事情开始 - 向量化循环，

假设我们有预先生成的索引数组i

、

j

和

k

。

def cosang1(A, i, j, k):
    d1 = A[i] - A[j]
    d2 = A[k] - A[j]
    d1_hat = np.linalg.norm(d1, axis=1, keepdims=True)
    d2_hat = np.linalg.norm(d2, axis=1, keepdims=True)
    # someone will almost certainly suggest a better way to do this
    ang = np.einsum("ij, ij -> i", d1/d1_hat, d2/d2_hat)
    return ang

这将问题简化为计算索引数组，假设索引数组只占总计算时间的一小部分。我看不出有什么方法可以在不做这种事情的情况下避免冗余计算。

然后，如果我们愿意允许冗余计算，生成索引的最简单方法是使用

np.meshgrid

。

def cosang2(A):
    i = np.arange(len(A))
    i, j, k = np.meshgrid(i, i, i)
    i, j, k = i.ravel(), j.ravel(), k.ravel()
    return cosang1(A, i, j, k)

对于 Colab 上形状 (30, 3) 的

A

，Python 循环方法花费了 160ms，而这个解决方案花费了 7ms。

如果允许使用 Numba，则可以非常轻松地快速生成唯一的索引集。这基本上只是你的代码分解成一个函数：

from numba import jit
# generate unique tuples of indices of vectors
@jit(nopython=True)
def get_ijks(M):
    ijks = []
    for i in range(M):
        for j in range(M):
            for k in range(i+1, M):
                if j not in (i, k):
                    ijks.append((i, j, k))
    return ijks

（当然，我们也可以在整个循环中使用 Numba。）

这比冗余矢量化解决方案花费的时间略少于一半。

使用纯 NumPy

可能可以有效地生成索引。最初，我认为这会很简单：

i = np.arange(M)
j, k = np.triu_indices(M, 1)
i, j, k = np.broadcast_arrays(i, j[:, None], k[:, None])
i, j, k = i.ravel(), j.ravel(), k.ravel()

这不太正确，但也许可以以此为起点，提出一个 O(M) 循环来修复这些索引。现在没有时间了。

---

所有代码放在一起，以防您想运行它：

from numba import jit
import numpy as np
rng = np.random.default_rng(23942342)

M = 30
N = 3
A = rng.random((M, N))

# generate unique tuples of indices of vectors
@jit(nopython=True)
def get_ijks(M):
    ijks = []
    for i in range(M):
        for j in range(M):
            for k in range(i+1, M):
                if j not in (i, k):
                    ijks.append((i, j, k))
    return ijks

# proposed method 
def cosang1(A, i, j, k):
    d1 = A[i] - A[j]
    d2 = A[k] - A[j]
    d1_hat = np.linalg.norm(d1, axis=1, keepdims=True)
    d2_hat = np.linalg.norm(d2, axis=1, keepdims=True)
    ang = np.einsum("ij, ij -> i", d1/d1_hat, d2/d2_hat)
    return ang

# another naive implementation
def cosang2(A):
    i = np.arange(len(A))
    i, j, k = np.meshgrid(i, i, i)
    i, j, k = i.ravel(), j.ravel(), k.ravel()
    return cosang1(A, i, j, k)

# naive implementation provided by OP
def cosang0(A):
    angles = []
    for i in range(len(A)):
        for j in range(len(A)):
            for k in range(i+1, len(A)):
                if j not in (i,k):
                    d1 = A[i] - A[j]
                    d2 = A[k] - A[j]
                    ang = np.dot(d1/np.linalg.norm(d1), d2/np.linalg.norm(d2))
                    angles.append(ang)
    return angles

%timeit cosang0(A)

%timeit get_ijks(len(A))
ijks = np.asarray(get_ijks(M)).T
%timeit cosang1(A, *ijks)

%timeit cosang2(A)

# 180 ms ± 34.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# 840 µs ± 68.3 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
# 2.19 ms ± 126 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# <ipython-input-1-d2a3835710f2>:26: RuntimeWarning: invalid value encountered in divide
#   ang = np.einsum("ij, ij -> i", d1/d1_hat, d2/d2_hat)
# 8.13 ms ± 1.78 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

cosangs0 = cosang0(A)
cosangs1 = cosang1(A, *ijks)
cosangs2 = cosang2(A)
np.testing.assert_allclose(cosangs1, cosangs0)  # passes