python - 二维中一组 N 个点之间的有向距离

Question

我在平面上有 N 个点存储为 (N, 2) 形 numpy 数组，我需要计算它们之间的力:F_ij = V(|r_i - r_j|)(r_i - r_j)，其中 V 是仅取决于点 i 和 j 之间的距离的函数。在代码语言中，我需要更有效地多次运行以下命令:

import numpy as np
def V(dist):  #Here, dist is a float
    return np.exp(-dist)/dist


N = 10000
r = np.random.rand((N, 2))
F = np.zeros((N,N,2))

for i in range(N):
    for j in range(N):
        r_rel = r[i, :] - r[j, :]  #Relative position
        dist_rel = np.sqrt(r_rel[:, 0]**2 + r_rel[:, 1]**2)
        F[i, j, :] = V(dist_rel)*r_rel

这也许可以通过使用一些内存存储技巧(这个 3-tensor 是内存昂贵的)来更有效地执行，通过只执行一半的 double-for 操作(因为 F[i,j,:]=-F [i,j,:]), 等等。所以，我问如何做到这一点。

scipy 模块的函数 cdist 与此非常相似，运行速度比上面的代码快 x12 倍，所以，一方面，也许在 scipy/numpy 中有一个函数可以做我想做的事情，另一方面，上面的代码可以写得更有效率。

谢谢你的帮助

Answer 1

我想到了两种方法。

首先，是使用scipy cdist .但它有几个问题。因为，它期望自定义距离函数的输出是一个标量，所以我们需要分别计算力的 x 和 y 坐标。其次，它将再次对所有对 i,j 和 j,i 进行计算。所以，正如问题中提到的，我们可以将计算减少到 n^2/2。在 N = 1000 时，需要 14.141 秒。对于 N = 10000，需要永远。在 8 GB、8 核 mac 上进行的计算。

import time
import numpy as np 
from scipy.spatial import distance

N = 1000
r = np.random.rand(N, 2)

def force(a, b):
    def V(dist):  #Here, dist is a float
        return np.exp(-dist)/(dist+1)

    r_rel = a - b  #Relative position
    dist_rel = np.sqrt(r_rel[0]**2 + r_rel[1]**2)
    return V(dist_rel)*r_rel

def force_x(a, b):
    return force(a,b)[0]
def force_y(a, b):
    return force(a, b)[1]

t1 = time.time()
# Calculate x and y coordinates separately
F_x = distance.cdist(r, r, metric=force_x)
F_y = distance.cdist(r, r, metric=force_y)
print("Time taken is = ", time.time() - t1) # takes 14.141s

其次，负责 n^2/2 优化的方法是使用 n^2 循环，但使用 numba 来加速计算。如果你有一个 gpu，你可以通过 numba vectorize 进一步加速。 在 N = 1000 时，需要 0.240 秒。对于 N = 10000，需要大约 25 秒。

import numba as nb
@nb.jit(nb.float64[:,:,:](nb.float64[:,:]), nopython=True, cache=True)
def distance(r):
    N = r.shape[0]
    F = np.zeros((N,N,2))
    for i in range(N):
        for j in range(i):
            r_rel = r[i] - r[j]  #Relative position
            dist_rel = np.sqrt(r_rel[0] ** 2 + r_rel[1] ** 2)
            V = np.exp(-dist_rel) / (dist_rel+1)    
            F[i,j] = V * r_rel
            F[j,i] = -1 * F[i,j]
    return F

t1 = time.time()
F = distance(r)
print("Time taken is = ", time.time() - t1) # takes 0.240s

因此，除了在评论中的链接中制作 O(n logn) 方法外，numba 方法似乎在可接受的时间内适用于问题中的示例大小。