在我的情况下如何使循环更有效？ （麻木）

Question

我正在运行 python 代码来计算某些坐标之间的距离。 原始数据如下：

a = np.array([[1,40,70],[2,41,71],[3,42,73]])    #id, latitude, longitude

我期待得到每对之间的距离，结果应该如下所示：

[1, 2, 100(km)]
[1, 3, 200(km)].
[2, 1, 100(km)]
[2, 3, 300(km)]
[3, 1, 200(km)]
[3, 2, 300(km)]

结果应包含 pair(m,n) 和 pair(n,m)

实际数据有39000列，所以我对代码效率有很大的要求。 目前我正在使用一个非常愚蠢的双循环：

line = 0
result = np.zeros((6,3))
for i in a:
    for j in a:
         dis = getDistance(i[1],i[2],j[1],j[2])  # this is the function i made to calculate distance between two coordinates
         result[line] = [i[0],j[0],dis]
         line += 1

谁能帮我改进代码？

Answer 1

如果你的距离 function 很简单，可以尝试找到distance_matrix并将其转换为类似于np.enumerate(distance_matrix)的 output 的数据结构：

def get_dist(X, Y):
    return 100*np.hypot(X-X[:,None], Y-Y[:,None])

M = np.array([[1,40,70],[2,41,71],[3,42,73]])
names, X, Y = np.transpose(M)
distance_matrix = get_dist(X, Y)
>>> list(np.enumerate(distance_matrix))
[((0, 0), 0.0),
 ((0, 1), 141.4213562373095),
 ((0, 2), 360.5551275463989),
 ((1, 0), 141.4213562373095),
 ((1, 1), 0.0),
 ((1, 2), 223.60679774997897),
 ((2, 0), 360.5551275463989),
 ((2, 1), 223.60679774997897),
 ((2, 2), 0.0)]

请注意，我们需要为索引添加不同的名称。 此外，这些名称和距离的类型不同，因此我们不能将它们都保存在一个（非结构化）数组中。 您可能希望避免迭代np.ndeumerate并以不同的方式查找名称：

x,y = np.indices([len(M), len(M)])
>>> names[x].ravel(), names[y].ravel(), distance_matrix.ravel()
(array([1, 1, 1, 2, 2, 2, 3, 3, 3]),
 array([1, 2, 3, 1, 2, 3, 1, 2, 3]),
 array([  0.        , 141.42135624, 360.55512755, 141.42135624,
          0.        , 223.60679775, 360.55512755, 223.60679775,
          0.        ]))

或者：

>>> np.transpose([names[x].ravel(), names[y].ravel(), dist_matrix.ravel()])
array([[  1.        ,   1.        ,   0.        ],
       [  1.        ,   2.        , 141.42135624],
       [  1.        ,   3.        , 360.55512755],
       [  2.        ,   1.        , 141.42135624],
       [  2.        ,   2.        ,   0.        ],
       [  2.        ,   3.        , 223.60679775],
       [  3.        ,   1.        , 360.55512755],
       [  3.        ,   2.        , 223.60679775],
       [  3.        ,   3.        ,   0.        ]])

Answer 2

首先，您要进行两次计算。 你的代码：

for i in a:
    for j in a:

计算从i到j和从j到i的距离，即使它们是相同的距离。 您只需执行以下操作即可将时间减半

for i in range(len(a)):
   for j in range(i+1, len(a)):
      D[i,j] = distance_calc(i,j)

即使您需要两个方向的距离，我也不会计算两次，只需在两个位置分配值即可。 但是如果你可以矢量化你的代码，它可能会加速

for i in range(len(a)):
   D[i,i+1:] = distance_calc(i)

为您的问题举一个例子，任何计算都将以弧度进行，所以我想首先将您的位置转换为弧度而不是度数。 即便如此，我们也可以加快速度。 我在这里假设您正在使用 Haversine estiamte 来计算球形地球上的距离。 我的算法基于您对问题的评论中引用的这篇文章中的信息。 首先，我将考虑N=1000点而不是你的全套。

import numpy as np
import time

r_Earth = 6371
N = 1000

def haversine_one_element(data1, data2):
    # calculates the Haversine distance one element at a time
    lat1 = data1[0]                     
    lng1 = data1[1]         

    lat2 = data2[0]                     
    lng2 = data2[1]         

    diff_lat = lat1 - lat2
    diff_lng = lng1 - lng2
    d = np.sin(diff_lat/2)**2 + np.cos(lat1)*np.cos(lat2) * np.sin(diff_lng/2)**2
    return 2 * r_Earth * np.arcsin(np.sqrt(d))            

def haversine_slow(a, N, do_all = True):
    D = np.zeros((N,N))
    for i in range(N):
        if do_all:
            for j in range(N):
                D[i,j] = haversine_one_element(a[i,1:], a[j,1:])
        else: 
            for j in range(i+1,N): # note that D[i,i] = 0 so we can skip it
                D[i,j] = D[j,i] = haversine_one_element(a[i,1:], a[j,1:])

a = np.array([np.arange(N).ravel(),
              np.random.random(N) * 360 - 180,
              np.arcsin(np.random.random(N) * 2 - 1) * 180 / np.pi]).transpose()
a[:,1:] = a[:,1:] * np.pi / 180

# Doing the entire array
start = time.time()
haversine_slow(a,N)
print(time.time() - start) # 8.777195453643799

# Doing only the upper half
start = time.time()
haversine_slow(a,N,do_all = False)
print(time.time() - start) # 4.547634840011597

但是，我们可以重写我们的 Haversine 公式以一次获取大量值，以便我们可以一次计算从一个点到所有其他点的距离。

def haversine(data1, data2):
    # data1, data2 are the data arrays with 2 cols and they hold
    # lat., lng. values in those cols respectively
    lat1 = data1[0]                     
    lng1 = data1[1]         

    lat2 = data2[0,:]                     
    lng2 = data2[1,:]         

    diff_lat = lat1 - lat2
    diff_lng = lng1 - lng2
    d = np.sin(diff_lat/2)**2 + np.cos(lat1)*np.cos(lat2) * np.sin(diff_lng/2)**2
    return 2 * r_Earth * np.arcsin(np.sqrt(d))            

def haversine_d(a,N):
    D = np.zeros((N,N))
    for i in range(N):
        d1 = a[i,1:]
        d2 = a[i+1:,1:].transpose()
        D[i,i+1:] = haversine(d1, d2)
        D[i+1:,i] = D[i, i+1:].transpose()
    # only return flattened upper triangular part of matirx as 
    return(D)

start = time.time()
d1 = haversine_d(a,N)
print(time.time() - start) # 0.03420424461364746

因此，对于我这里的情况，计算时间又减少了 100 倍，我认为对于像你所说的那样更长的向量，它会减少更多。 （受内存限制问题的影响，见下文。）

然而，我认为我们还可以做一件事，如果我们用线性代数而不是三角函数来看待这个问题。 如果你学过线性代数，你应该知道点积和余弦之间的关系，我们可以通过将位置描述为(x,y,z)中的向量而不是经度和纬度来利用这些信息来发挥我们的优势：

def get_xyz(a):
    x = np.cos(a[:,1]) * np.cos(a[:,2])
    y = np.cos(a[:,1]) * np.sin(a[:,2])
    z = np.sin(a[:,1]) 
    v = np.stack([x, y, z]).transpose()
    return(np.hstack([a[:,:1], v]))

def distance_dot(a,N):
    b = get_xyz(a)
    
    D = np.zeros((N,N))
    for i in range(N):
        D[i,i+1:] = np.arccos(b[i,1] * b[i+1:,1] + \
                              b[i,2] * b[i+1:,2] + \
                              b[i,3] * b[i+1:,3]) * r_Earth
        D[i+1:,i] = D[i, i+1:].transpose()
    return(D)

start = time.time()
distance_dot(a,N) # 0.019799470901489258
print(time.time() - start)

这种代数方法应该在浮点计算的截断限制内给出相同的答案； 如果您尝试使用==比较 Haversine 和线性代数的测试，它将失败，但如果您执行<= 10**-10之类的操作，它将通过。

这些方法在我的 1000 点数据集上将计算时间减少了大约 400 倍。 对于更大的数据集，我没有在N=40000的情况下运行前两个，但对于后两个，Haversine 和线性代数测试几乎同时出现（290 秒），但这可能是因为问题最终需要所有我笔记本电脑上的 memory。 如果 memory 不是问题，我预计问题会大致按N**2缩放，因此 Haversine 大约为 60 s，线性代数为 30 s。

如果您开始遇到 memory 问题，可能有一种方法可以通过计算数组的块并像 go 一样保存到磁盘来加速它，但暂时需要解决这个问题，除非您需要生成生产级代码或一遍又一遍地这样做我可能不会打扰。