Python - 低效的空间距离计算（如何加快速度）答案

【问题标题】：Python - inefficient spatial distance calculation (how can it be speed up)Python - 低效的空间距离计算（如何加快速度）
【发布时间】：2018-07-15 04:16:58
【问题描述】：

我目前正在尝试在 Python 中进行一些地理编码。该过程如下：我有两个具有纬度和经度值的数据框（df1 和 df2，房屋和学校），并且希望为 df1 中的每个观察结果在 df2 中找到最近的邻居。我使用以下代码：

from tqdm import tqdm
import numpy as np
import pandas as pd
import math 

def distance(lat1, long1, lat2, long2):
        R = 6371 # Earth Radius in Km
        dLat = math.radians(lat2 - lat1) # Convert Degrees 2 Radians 
        dLong = math.radians(long2 - long1)
        lat1 = math.radians(lat1)
        lat2 = math.radians(lat2)
        a = math.sin(dLat/2) * math.sin(dLat/2) + math.sin(dLong/2) * math.sin(dLong/2) * math.cos(lat1) * math.cos(lat2)
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
        d = R * c
        return d

dists = []
schools =[]
for index, row1 in tqdm(df1.iterrows()):
    for index, row2 in df2.iterrows():
        dists.append(distance(row1.lat, row1.lng, row2.Latitude, row2.Longitude))
    schools.append(min(dists))
    del dists [:]

df1["school"] = pd.Series(schools)

代码有效，但需要很长时间。使用 tqdm，我得到 df1 每秒 2 次迭代的平均速度。作为比较，我在 STATA 中使用 geonear 完成了整个任务，df1 (950) 中的所有观察都需要 1 秒。我在 geonear 的帮助文件中读到他们使用聚类，不计算所有距离，而只计算最近的距离。但是，在我添加一个集群功能（这也可能需要 CPU 能力）之前，我想知道是否有人看到了一些方法来加快进程（我是 python 新手，可能有一些效率低下的代码会减慢进程） .或者是否有一个包可以更快地完成这个过程？

如果它需要比 STATA 更长的时间，我会没事的，但不会接近 7 分钟......

提前谢谢你

【问题讨论】：

标签： python algorithm performance geospatial distance

【解决方案1】：

您执行此操作的方式很慢，因为您使用的是 O(n²) 算法：每行查看其他行。 Georgy's answer 虽然引入了矢量化，但并没有解决这个根本的低效率问题。

我建议将您的数据点加载到 kd-tree 中：此数据结构提供了一种快速查找多个维度中最近邻居的方法。这样一棵树的构建耗时O(n log n)，查询耗时O(log n)，所以总时间O(n log n ).

如果您的数据本地化到可以通过平面很好地近似的地理区域，请投影您的数据，然后在二维中执行查找。否则，如果您的数据是全球分散的，请投影到 spherical cartesian coordinates 并在那里执行查找。

您可以如何执行此操作的示例如下所示：

#/usr/bin/env python3

import numpy as np
import scipy as sp
import scipy.spatial

Rearth = 6371

#Generate uniformly-distributed lon-lat points on a sphere
#See: http://mathworld.wolfram.com/SpherePointPicking.html
def GenerateUniformSpherical(num):
  #Generate random variates
  pts      = np.random.uniform(low=0, high=1, size=(num,2))
  #Convert to sphere space
  pts[:,0] = 2*np.pi*pts[:,0]          #0-360 degrees
  pts[:,1] = np.arccos(2*pts[:,1]-1)   #0-180 degrees
  #Convert to degrees
  pts = np.degrees(pts)
  #Shift ranges to lon-lat
  pts[:,0] -= 180
  pts[:,1] -= 90
  return pts

def ConvertToXYZ(lonlat):
  theta  = np.radians(lonlat[:,0])+np.pi
  phi    = np.radians(lonlat[:,1])+np.pi/2
  x      = Rearth*np.cos(theta)*np.sin(phi)
  y      = Rearth*np.sin(theta)*np.sin(phi)
  z      = Rearth*np.cos(phi)
  return np.transpose(np.vstack((x,y,z)))

#For each entry in qpts, find the nearest point in the kdtree
def GetNearestNeighbours(qpts,kdtree):
  pts3d        = ConvertToXYZ(qpts)
  #See: https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.spatial.KDTree.query.html#scipy.spatial.KDTree.query
  #p=2 implies Euclidean distance, eps=0 implies no approximation (slower)
  return kdtree.query(pts3d,p=2,eps=0) 

#Generate uniformly-distributed test points on a sphere. Note that you'll want
#to find a way to extract your pandas columns into an array of width=2, height=N
#to match this format.
df1 = GenerateUniformSpherical(10000)
df2 = GenerateUniformSpherical(10000)

#Convert df2 into XYZ coordinates. WARNING! Do not alter df2_3d or kdtree will
#malfunction!
df2_3d = ConvertToXYZ(df2)
#Build a kd-tree from df2_3D
kdtree = sp.spatial.KDTree(df2_3d, leafsize=10) #Stick points in kd-tree for fast look-up

#Return the distance to, and index of, each of df1's nearest neighbour points
distance, indices = GetNearestNeighbours(df1,kdtree)

【讨论】：

@Georgy：Python 中的向量化通常是获得良好性能的关键组成部分。您的代码以一种简单的 OP 方式利用了这种性能来源。对于较小的数据集，我可能会很好地工作，特别是因为它的构建成本比 kd-tree 小。不过，在某些时候，相对时间复杂性将占主导地位，大猩猩会出击:-)
@Georgy：我上传了一些代码，展示了如何使用 kd-tree。
@Richard 我尝试了 Georgy 的解决方案，因为它看起来有点短。但是，如上所述，我可能不得不回到您的解决方案，因为我的实际数据集包含超过 300 万个观察值。提前谢谢你，也为链接。对于这样的高效编程，您可能还有更多要阅读的内容吗？
@user27074：我建议多读几遍 Skiena 的《算法设计手册》的前半部分。这是一个很好的算法介绍。之后，只要学习所有你能学到的算法。

【解决方案2】：

pandas 提高效率的关键是对整个数据帧/系列执行操作，而不是逐行执行。所以，让我们这样做吧。

for index, row1 in tqdm(df1.iterrows()):
    for index, row2 in df2.iterrows():

在这里计算两个数据帧的笛卡尔积。这可以像这样更快地完成：

df_product = pd.merge(df1.assign(key=0, index=df1.index), 
                      df2.assign(key=0), 
                      on='key').drop('key', axis=1)

（代码取自here）。我还添加了一个索引为df1 的列，我们稍后将需要它来计算df1 中每个实体的距离min。

现在使用 numpy 以矢量化方式计算所有增量、以弧度为单位的纬度、a、c 和距离：

dLat = np.radians(df_product['Latitude'] - df_product['lat'])
dLong = np.radians(df_product['Longitude'] - df_product['lng'])
lat1 = np.radians(df_product['lat'])
lat2 = np.radians(df_product['Latitude'])
a = (np.sin(dLat / 2) ** 2 
     + (np.sin(dLong / 2) ** 2) * np.cos(lat1) * np.cos(lat2))
c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1 - a))
df_product['d'] = R * c

现在，从df_product 开始，我们只留下我们之前添加的索引列和距离列。我们按索引对距离进行分组，计算相应的最小值并将它们分配给df1['schools']，就像您在代码中所做的那样。

df1['schools'] = df_product.loc[:, ['index', 'd']].groupby('index', axis=0).min()

就是这样。对于每个数据帧中的 1000 行，对我来说，一切都只需要不到一秒钟的时间。

【讨论】：

@Richard 是的，好吧.. 这么多行会出现内存错误。
@Richard 我不知道在这种情况下避免MemoryError 的最佳方法。要么是按块处理数据，要么完全避免使用 pandas 并使用生成器。我自己以前从未遇到过这样的问题。所以，我希望 OP 对我的解决方案感到满意。
pandas 应该能够轻松处理一百万行，尽管它可能会在合并时阻塞，因为这可能会在内存或时间上产生 O(N²) 扩展在熊猫中实施不佳。我想间接指出的是，您的解决方案仍在 O(N²) 时间内运行，只是比 OP 更有效。这意味着无论内存限制如何，它的扩展性都很差。
@Georgy 非常感谢！我使用了这段代码，确实更快。但是，我的想法是在子样本（950 所房屋 x 5300 所学校）上编写代码，然后将其用于我的实际数据（300 万所房屋 x 10000 所学校）。所以我稍后也会尝试Richard的解决方案（我仍然要清理其他数据集）。由于我主要通过在线资源学习 Python，这些资源主要针对应用程序而不是性能，我很好奇您是否有任何关于此主题的资源可供阅读（例如，为什么这比我原来的方法更快）。谢谢！
@user27074 我对 numpy 的了解都来自 Stack Overflow 和 numpy 文档。所以，我在这里真的不能给任何建议。但我知道有一本书叫《高性能Python》。也许它对你有任何帮助。