Python - 空间距离计算效率低下（如何提速）

Question

我目前正在尝试使用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后，每秒只能处理2个df1的迭代。相比之下，我使用STATA中的geonear完成了整个任务，对于df1中所有观察值（950个），仅需1秒钟。我在geonear的帮助文件中读到，他们使用聚类来避免计算所有距离，而只计算最近的距离。然而，在添加聚类功能之前（这也可能会消耗CPU资源），我想知道是否有人看到了一些可以加快当前进程速度的方法（我是Python的新手，可能存在一些低效的代码导致进程变慢）。或者，也许有一个可以更快地处理这个任务的软件包？

如果它需要的时间比STATA长，那么也没关系，但不要接近7分钟...

先感谢您的帮助。

- user27074

Answer 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)

- Richard

---

@Georgy：Python中的向量化通常是优秀性能的关键组成部分。你的代码以一种简单的方式利用了这种性能源，对于OP来说它很可能在较小的数据集上非常有效，特别是因为它的构建成本比kd树更低。不过到了某个时候，相对时间复杂度将会占主导地位，大猩猩就会出现问题了 :-) - Richard

@Georgy：我上传了一些代码，展示了如何使用kd树。 - Richard

@Richard 我尝试了Georgy的解决方案，因为它看起来比较简短。然而，如上所述，由于我的实际数据集包含超过3百万个观测值，我可能需要回到你的解决方案。提前感谢您，也感谢提供的链接。您是否有更多关于高效编程的阅读材料？ - user27074

@user27074：我建议多次阅读斯基纳的《算法设计手册》的前半部分。它是算法入门的好材料。之后，尽可能学习所有的算法。 - Richard

Answer 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)

（代码摘自这里）我还添加了一列用于索引df1的索引，我们稍后需要用它来计算每个实体到df1的距离的最小值。

---

现在使用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行，我所需的时间不到一秒钟。

- Georgy

---

@Richard 是的，对于这么多行会出现内存错误。 - Georgy

@Richard 我不知道在这种情况下避免MemoryError的最佳方法。要么是按块处理数据，要么是完全避免使用pandas并使用生成器。我以前从未遇到过这样的问题。因此，我希望OP对我的解决方案感到满意。 - Georgy

Pandas应该能够轻松处理一百万行数据，尽管在合并时可能会出现问题，因为如果在Pandas中实现不好的话，它可能会导致内存或时间上的_O(N²)_扩展。我想间接指出的是，你的解决方案仍然在_O(N²)_时间内运行，只是比OP的更有效率。这意味着无论内存限制如何，它都将扩展得很差。 - Richard

@Georgy 非常感谢您！我使用了这段代码，确实快得多。然而，我的想法是在一个子样本（950个房屋x5300所学校）上编程代码，然后将其用于我的实际数据（3百万个房屋x1万个学校）。所以我稍后也会尝试Richard的解决方案（我还需要清理另一个数据集）。由于我主要通过在线资源学习Python，这些资源大多数是面向应用而不是性能的，所以我很好奇您是否有任何阅读关于这个主题的来源（例如，为什么这比我的原始方法更快）。谢谢！ - user27074

@user27074，我学习NumPy的大部分知识都来自于Stack Overflow和NumPy文档。所以在这里我无法给出任何建议。但是我知道有一本书叫做《高性能Python》，也许它可以对你有所帮助。 - Georgy