如何加快 Pandas 地理距离计算？

Question

I am a Pandas / Python Beginner and I don't know how I can speed up my code. I have a unsorted Pandas Dataframe called test with about 10.000 rows with multiple columns, including latitude and longitude. I want to know, for each row, how many other rows are close by (within distance of a threshold ie 10 km).

I tried doing that:

import numpy as np
def distance(lat1, lon1, lat2, lon2, to_radians=True, earth_radius=6371):

    lat1, lon1, lat2, lon2 = np.radians([lat1, lon1, lat2, lon2])

    a = np.sin((lat2-lat1)/2.0)**2 + \
        np.cos(lat1) * np.cos(lat2) * np.sin((lon2-lon1)/2.0)**2

    return earth_radius * 2 * np.arcsin(np.sqrt(a))

and:

["Number_of_neighbours"] = 0
distance = 5
i = 0
y = 0 
for i in range(0, len(test)):
 for y in range(0,len(test)):
 x = haversine(lat1 = test['latitude'].loc[i] , lon1 = test['longitude'].loc[i] , lat2 = test['latitude'].loc[y] , lon2 = test['longitude'].loc[y])
 
 if x <= distance and x != 0:  
  test.at[i,'Number_of_neighbours']= 1 + test.loc[i, 'Number_of_neighbours']
  

But Jupyter takes for ever to compute the result. Do you have any suggestion or a more performant solution in mind? Thank you very much in advance!

Answer 1

Since your data isn't too big, you can rewrite your distance function to take advantage of numpy's broadcasting

def haversine(lat1, lon1, lat2, lon2, to_radians=True, earth_radius=6371):
    lat1, lon1, lat2, lon2 = np.radians([lat1, lon1, lat2, lon2])

    a = np.sin((lat2-lat1[:,None])/2.0)**2 + \
        np.cos(lat2) * np.cos(lat1[:,None]) * np.sin((lon2[:,None]-lon1)/2.0)**2

    return earth_radius * 2 * np.arcsin(np.sqrt(a))

distance = 5
dist = haversine(test['latitude'],test['longitude'], test['latitude'], test['longitude'])

test['neighbors'] = (dist < distance).sum(-1)

All that took about 7s on my system on 10000 rows.