加速Pandas过滤

Question

我有一个包含37456153行和3列的Pandas数据框，这些列分别是：[时间戳, 跨度, 高度]。每个时间戳大约有62000行的跨度和高度数据。比如，当我筛选时间戳 = 17210时，数据看起来是这样的：

        Timestamp       Span  Elevation
94614       17210  -0.019766     36.571
94615       17210  -0.019656     36.453
94616       17210  -0.019447     36.506
94617       17210  -0.018810     36.507
94618       17210  -0.017883     36.502

...           ...        ...        ...
157188      17210  91.004000     33.493
157189      17210  91.005000     33.501
157190      17210  91.010000     33.497
157191      17210  91.012000     33.500
157192      17210  91.013000     33.503

从上面可以看到，跨度的数据并不是均匀分布的，而我其实需要它是均匀的。所以我想出了以下代码，把它转换成均匀分布的格式。我知道我想分析的开始和结束位置。然后我定义了一个增量参数，作为我的步长。我创建了一个叫mesh的numpy数组，里面存放我想要的均匀分布的跨度数据。最后，我决定对数据框中的某个时间戳（代码中的17300）进行迭代，测试一下它的运行速度。代码中的for循环计算每个增量范围内的平均高度值，范围是±0.5增量。

我的问题是：过滤数据框并计算平均高度的时间为603毫秒，这只是单次迭代的时间。对于给定的参数，我需要进行9101次迭代，这样一来，这个循环大约需要1.5小时才能完成。而且，这只是针对一个时间戳的值，而我有600个（做完所有的需要900小时？！）。

有没有什么办法可以加快这个循环的速度？非常感谢任何建议！

# MESH GENERATION
start = 0
end = 91
delta = 0.01

mesh = np.linspace(start,end, num=(end/delta + 1))
elevation_list =[]

#Loop below will take forever to run, any idea about how to optimize it?!

for current_loc in mesh:
    average_elevation = np.average(df[(df.Timestamp == 17300) & 
                                      (df.Span > current_loc - delta/2) & 
                                      (df.Span < current_loc + delta/2)].Span)
     elevation_list.append(average_elevation)

Answer 1

这里有个想法——可能还是太慢了，但我想分享一下。首先，准备一些假数据。

df = pd.DataFrame(data={'Timestamp': 17210, 
                        'Span': np.linspace(-1, 92, num=60000), 
                        'Elevation': np.linspace(33., 37., num=60000)})

接下来，把你创建的网格数组转换成一个数据框（dataframe），并添加一个偏移的条目，这样数据框中的每个条目就代表新均匀间隔的一步。

mesh_df = pd.DataFrame(mesh, columns=['Equal_Span'])
mesh_df['Equal_Span_Prev'] = mesh_df['Equal_Span'].shift(1)
mesh_df = mesh_df.dropna()

然后，我想把这个数据框和更大的数据集结合起来，条件是这个条目要在两个 Equal_Span 列之间。虽然在pandas中可能有办法做到，但在SQL中表达笛卡尔类型的连接似乎更简单，所以我先把所有数据放到一个内存中的sqlite数据库里。如果你遇到内存问题，我建议把它做成基于文件的数据库。

import sqlite3
con = sqlite3.connect(':memory:')
df.to_sql('df', con, index=False)
mesh_df.to_sql('mesh_df', con, index=False)

这是主要的查询。我的测试数据大约花了1分30秒，所以在完整数据集上可能还是会很耗时。

join_df = pd.read_sql("""SELECT a.Timestamp, a.Span, a.Elevation, b.Equal_Span
                         FROM df a, mesh_df b
                         WHERE a.Span BETWEEN b.Equal_Span_Prev AND b.Equal_Span""", con)

不过，一旦数据变成这个形式，获取想要的平均值就变得简单又快速了。

join_df.groupby(['Timestamp','Equal_Span'])['Elevation'].mean()

Answer 2

你可以使用 np.searchsorted 来把整个过程变得更简单高效。我对 pandas 不是很熟悉，但类似下面的代码应该可以工作，而且在我的系统上运行得相当快。使用 chrisb 的示例数据：

In [8]: %%timeit
   ...: mesh = np.linspace(start, end, num=(end/delta + 1))
   ...: midpoints = (mesh[:-1] + mesh[1:]) / 2
   ...: idx = np.searchsorted(midpoints, df.Span)
   ...: averages = np.bincount(idx, weights=df.Elevation, minlength=len(mesh))
   ...: averages /= np.bincount(idx, minlength=len(mesh))
   ...: 
100 loops, best of 3: 5.62 ms per loop  

这个方法比你的代码快了大约 3500 倍：

In [12]: %%timeit
    ...: mesh = np.linspace(start, end, num=(end/delta + 1))
    ...: elevation_list =[]
    ...: for current_loc in mesh:
    ...:     average_elevation = np.average(df[(df.Span > current_loc - delta/2) & 
    ...:                                       (df.Span < current_loc + delta/2)].Span)
    ...:     elevation_list.append(average_elevation)
    ...: 
1 loops, best of 3: 19.1 s per loop

---

编辑 那这个是怎么工作的呢？在 midpoints 中，我们存储了一个有序的边界列表，这些边界用来划分不同的区间。接着，我们在这个有序列表上使用 searchsorted 进行二分查找，得到 idx，它基本上告诉我们每个数据点属于哪个区间。接下来只需要把每个区间里的值进行分组，这就是 bincount 的作用。给定一个整数数组，它会统计每个数字出现的次数。如果给定一个整数数组和一个对应的 weights 数组，它就不会简单地在区间计数上加 1，而是会加上 weights 中对应的值。通过两次调用 bincount，你可以得到每个区间的总和和项目数量：把它们相除就能得到区间的平均值。