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张量会占用大量内存），或者仅执行一半的double-for操作（因为F [i，j，：] =-F [i，j ,:]），等等。我要求这样做。

scipy模块的函数cdist与此类似，并且运行速度比上述代码快12倍，因此，一方面，也许有一个函数可以执行scipy / numpy中的功能，另一方面，上述代码可以更有效地编写。

感谢您的帮助

Answer 1

我想到两种方法。

首先，要使用样式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方法似乎还可以在可接受的时间内处理问题中的示例大小。

Answer 2

这应该是对@Deepak Sainis答案的评论。除了他已经很好的答案之外，还有一些性能关键问题。

1。数组声明

nb.float64[:,:,:](nb.float64[:,:])

这意味着输入和输出数组在内存中都可以是不连续的。这通常会阻止SIMD向量化。如果数组是绝对连续的，则可以声明

nb.float64[:,:,::1](nb.float64[:,::1])

这意味着这些目标是C邻接的。

或者您可以简单地删除声明，Numba将正确检测数据类型和连续性本身。

2。使用许多可以展开的小循环

每个附加循环都会带来一些开销。在此示例中，已知数组r的第二个暗号将为2。您可以手动展开“隐藏循环”。在此示例中，这将性能提高了一个数量级以上。

#eg. this
r_rel = r[i] - r[j]
#can be rewritten to
r_rel_1 = r[i,0] - r[j,0]
r_rel_2 = r[i,1] - r[j,1]

3。仅计算一次结果并复制声音是个好主意

F[j,i] = -1 * F[i,j]

如果可能的话，此行阻止SIMD向量化，并且成本很高（连续读取/写入数组）。与直接计算所有值相比，此简单操作花费的时间将近一点。但是，如果您知道存在某种形式的对称性，则可以完全避免此复制过程。

4。如果仍然覆盖结果，请使用np.zeros

这大约是计算数组上半部分成本的（1/3）。在后端，总是有一个np.empty（）来分配内存。如果使用np.zeros，则存在一个嵌套循环，该循环将零写入整个数组，这非常昂贵。

5。并行化

对于小问题，这可能会对性能产生负面影响，但是较大的问题会从中受益。缓存不适用于并行函数，并且编译时间更长。

单线程版本

@nb.jit(fastmath=True,error_model='numpy',cache=True)
def distance(r):
    N = r.shape[0]
    F = np.empty((N,N,2))
    #give the compile all information available
    assert r.shape[1]==2
    for i in range(N):
        #change i to N to get all results
        for j in range(i):
            r_rel_1 = r[i,0] - r[j,0]
            r_rel_2 = r[i,1] - r[j,1]
            dist_rel = np.sqrt(r_rel_1 ** 2 + r_rel_2 ** 2)
            V = np.exp(-dist_rel) / (dist_rel+1)    
            F[i,j,0] = V * r_rel_1
            F[i,j,1] = V * r_rel_2
    return F

并行版本

@nb.njit(parallel=True,error_model='numpy',fastmath=True)
def distance(r):
    N = r.shape[0]
    F = np.empty((N,N,2))
    #give the compile all information available
    assert r.shape[1]==2
    for i in nb.prange(N):
        #change i to N to get all results
        for j in range(i):
            r_rel_1 = r[i,0] - r[j,0]
            r_rel_2 = r[i,1] - r[j,1]
            dist_rel = np.sqrt(r_rel_1 ** 2 + r_rel_2 ** 2)
            V = np.exp(-dist_rel) / (dist_rel+1)
            F[i,j,0] = V * r_rel_1
            F[i,j,1] = V * r_rel_2
    return F

时间

Core i7 Quadcoe第四代，Python 3.5，Numba 0.4dev 总是有第二个调用要测量的函数，以避免测量编译开销。仅计算了一半的数组。我评论了F[j,i] = -1 * F[i,j]行以进行直接比较。

N=1000    
@Deepak Saini: 109ms
single_threaded: 6.4ms
parallel:5ms

N=10000
@Deepak Saini: 10.43s
single_threaded: 0.61s
parallel:0.36s

Answer 3

没有理由在会导致np.exp下溢为零的距离上执行耗时的np.sqrt和np.exp(-dist)函数。大于零的最小双精度数为5e-324和-np.log(5e-324) = 744.44007192138122。现在，在您的示例中，最小距离仅为〜1.4，但我认为实际上该距离会更大。考虑到这一点，我们可以做到：

import numpy as np

def V(dist_sq, thr = 5e-324):  #Here, dist_sq is dist**2, a float
    d_sq_thr = np.log(thr)**2  
    mask = dist_sq < d_sq_thr  # mask is True where you want to do calculations
    out = np.sqrt(dist_sq[mask]) # square root only in place and only for valid distances
    out = np.exp(-out) / out  # same with `np.exp`
    return out, mask

i, j = np.triu_indices(N, 1)  # gives you the valid (i, j) values
r_rel = r[i] - r[j]
sq_distances = np.sum(r_rel**2, axis = 1)  # squared distances
v_, mask = V(sq_distances)
F[i[mask], j[mask], :] = v_ * r_rel[mask]
F[j[mask], i[mask], :] = -F[i[mask], j[mask], :]

现在，您可以将thr的值设置为大于机器精度的值，甚至可以加快速度，甚至可以按照@ max9111的建议开始将其包装在numba中。不是numba专家，但我认为这会起作用：

@nb.njit(parallel=True,error_model='numpy',fastmath=True)
def distance(r, thr = None):
    if thr == None:
        d_sq_thr = np.log(5e-324)**2
    else:
        d_sq_thr = np.log(thr)**2 
    N = r.shape[0]
    F = np.zeros((N,N,2))
    #give the compile all information available
    assert r.shape[1]==2
    for i in nb.prange(N):
        #change i to N to get all results
        for j in range(i):
            r_rel_1 = r[i,0] - r[j,0]
            r_rel_2 = r[i,1] - r[j,1]
            dist_rel_sq = r_rel_1 ** 2 + r_rel_2 ** 2
            if dist_rel_sq > d_sq_thr :
                continue
            dist_rel = np.sqrt(dist_rel_sq)
            V = np.exp(-dist_rel) / (dist_rel+1)
            F[i,j,0] = V * r_rel_1
            F[i,j,1] = V * r_rel_2
    return F

请记住，此方法在玩具问题上会比其他方法慢，因为距离检查只是开销（因为您的r值都限制在[0,1 ]），但是当使用较大距离时，这应该可以加快速度-甚至可以加快速度。

二维中的一组N点之间的定向距离

3 个答案: