python - Pandas 中更快的应用方法

Question

我有一个函数要应用于位置数据框。具体来说，我想附加一个新列，其中包含离每个站点最近的 10 个站点。以下似乎有效，但速度非常慢。

def distance(first_lat, first_lon, second_lat, second_lon):
    return ((first_lat - second_lat) ** 2 + (first_lon - second_lon) ** 2) ** 0.5


def load_site_list():
    '''
    This function generates a dataframe with all the available sites
    '''
    url = 'ftp://ftp.ncdc.noaa.gov/pub/data/noaa/isd-history.csv'
    cols = ["STATION NAME",
            "LAT",
            "LON"]
    df = pd.read_csv(url, parse_dates=False, usecols=cols)
    df = df.dropna(subset=['LAT'])
    df = df.dropna(subset=['LON'])
    df['LAT'] = df['LAT'].astype(float)
    df['LON'] = df['LON'].astype(float)
    return df

sites = load_site_list()
sites['closest'] = ""
for index, row in sites.iterrows():
    sites['dist'] = sites.apply(lambda line: distance(line['LAT'], line['LON'], row['LAT'], row['LON']), axis=1)
    sites.sort_values('dist', inplace=True)
    sites['closest'][index] = sites['STATION NAME'].iloc[1:11].tolist()

似乎 for 循环中生成与当前列的距离的第一行在每个循环中占用了一秒钟。这里有超过 10,000 行要循环...有更快的方法吗？

Answer 1

请注意，您的代码的时间复杂度为 O(n^2):在这种情况下，您要在 for 循环中的应用函数(即纯 Python)中计算 30k*30k=9 亿个距离。

pandas 中的矢量运算是用 C 语言实现的，因此如果您在单个矢量运算中计算所有距离，您将获得相对加速。

如果您有足够的 RAM，您可以进行笛卡尔连接，计算所有成对距离，然后进行排序、分组，然后取头部，如下所示:

# code to reduce memory usage
sites['site_code'] = pd.Categorical(sites['STATION NAME']).codes
sites['LAT'] = sites.LAT.astype(np.float16)
sites['LON'] = sites.LAT.astype(np.float16)
sites_small = sites[['site_code','LAT','LON']].copy()
sites_small.index = [0]*len(sites_small)

pairs = sites_small.join(sites_small,lsuffix='_x',rsuffix='_y')
pairs['dist'] = (pairs['LAT_x'] - pairs['LAT_y'])**2 + (pairs['LON_x'] - pairs['LON_y'])**2
pairs.sort_values(['STATION NAME_x','dist'], inplace = True) # actually, just sorting by dist is sufficient
pairs.groupby('STATION NAME_x').head(10)

不幸的是，您可能没有足够的 RAM:如果您将站点名称编码为 16 位整数，并将坐标编码为 16 位 float ，则每行需要 12 个字节(因为您要查看的是成对的) ，加上另外 8 个字节的索引(pandas 在连接中将这些带入 lonints；我不知道如何解决这个问题)，最终数据帧大约需要 20 字节 * 900m 行 = 18GB。在实践中可能更多，并且操作期间的内存使用峰值高于此(特别是排序将花费最长的时间，并且使用大量内存)。

我在我的机器上试过这个:我使用了大约 30GB，放弃等待完整排序，而是对 dist 小于 100 的子集进行排序。用时不到 5 分钟，大部分时间花在了连接上。

归根结底，您要进行近十亿次计算；如果您想以 C 的速度执行此操作而不必存储所有成对数据(pandas 中的直接方法就是这种情况)，您很可能必须使用 numpy 数组和/或使用 Cython 编写代码多处理。

更聪明的方法是避免进行十亿次计算，这涉及到知道哪些距离您不需要费心计算。这需要一些聪明的逻辑，但幸运的是，这是一个经过深入研究的 k-最近邻主题，它具有专门针对这种性质的问题设计的高效算法:

from sklearn.neighbors import NearestNeighbors
data = sites[['LAT','LON']].values
nbrs = NearestNeighbors(n_neighbors=10, algorithm='auto', metric = 'euclidean').fit(data)
distances, indices = nbrs.kneighbors(data)
indices

这需要不到一秒钟的时间来计算。恢复最近邻居的名称需要更长的时间:

df = pd.DataFrame(indices, index = sites['STATION NAME'].values)
df.replace(dict(enumerate(sites['STATION NAME'].values)), inplace = True)

(实际上，您可以通过使用带有一些堆叠/取消堆叠的 .merge() 方法来显着加快速度，但在这种情况下，它有点棘手，因为您的数据包含重复项。)