来自 groupby 的 Pandas 快速加权随机选择

Question

我有一个有趣的性能优化问题，目前是我们应用程序的瓶颈

鉴于 DataFrame 具有非唯一时间戳 index、id 和 weight 列（事件）和一系列时间戳（观察），我必须为每个观察分配一个随机事件 id，该事件发生在给定时间戳考虑权重。时间戳被限制到最接近的分钟，并且可以被视为从某个开始日期时间开始的分钟数。

测试数据生成：

import pandas as pd
import numpy as np
import random

from datetime import datetime as dt, timedelta as td

# typical date range is one month
start = dt(2020, 2, 1, 0, 0, 0)
end = dt(2020, 3, 1, 0, 0, 0)

# generate one event per minute
index = pd.date_range(start, end, freq='1min')
N = len(index)
events = pd.DataFrame({'id': np.arange(N), 'weight': np.random.random(N)}, index=index)

# generate some random events to simulate index duplicates
random_minutes = pd.to_datetime([start + td(minutes=random.randint(0, N)) for m in range(3*N)])
random_events = pd.DataFrame({'id': np.arange(3*N), 'weight': np.random.random(3*N)}, index=random_minutes)
events = pd.concat([events, random_events])

# observations, usually order or two orders of magnitude more records than events
observations = pd.Series([start + td(minutes=random.randint(0, N)) for m in range(10*N)])

样本数据点

>>> print(events.sort_index().to_string())
                     id    weight
2020-02-09 01:00:00   0  0.384927
2020-02-09 01:00:00  15  0.991314
2020-02-09 01:00:00  17  0.098999
2020-02-09 01:01:00   1  0.813859
2020-02-09 01:01:00   2  0.922601
2020-02-09 01:01:00   1  0.738795
2020-02-09 01:02:00   2  0.898842
2020-02-09 01:02:00  13  0.621904
2020-02-09 01:03:00  12  0.075857
2020-02-09 01:03:00   3  0.135762
2020-02-09 01:03:00   9  0.398885
...

>>> print(observations.sort_values().to_string())
12   2020-02-09 01:00:00
9    2020-02-09 01:00:00
44   2020-02-09 01:00:00
31   2020-02-09 01:01:00
53   2020-02-09 01:02:00
3    2020-02-09 01:02:00
6    2020-02-09 01:03:00

我目前最快的解决方案是通过索引groupby 事件，为每个记住样本的组函数返回。很难对其进行正确矢量化，因为每个组的许多记录可能会有所不同，并且我必须根据权重返回 ID。

%%timeit

from functools import partial

# create a per-minute random function returning id according to weights
randomizers = events.groupby(level=0).apply(
    lambda s: partial(
        np.random.choice, 
        s.id.values, 
        p=s.weight.values/s.weight.sum()
    )
)

# for each observation, find random generator and call it
selections = randomizers.loc[observations].apply(lambda f: f())

14.7 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

所以我的问题是，有没有一种更好、更快的方法来做我需要做的事情？我面临的主要问题：

1. 每分钟可以有多个事件，每个事件都有 ID 和概率
2. 每分钟的事件数是随机的，一分钟可以有 1 个，不同的可以有 20 个
3. 对于每个观察，我需要单独选择一个随机选项。

有什么想法吗？我正在考虑使用 numba，但也许有一些聪明的解决方案？

Answer 1

我可以建议您在此处获得性能的两点。

首先，访问groupby.apply 中的 id/weight 列会创建新系列，这很昂贵。如果您按日期对事件数据帧进行排序，则可以通过对原始 ndarray 进行切片来更有效地提取所需的输入。

另一点是关于 RNG。函数random.choice 是相当高级的，除了累积分布函数，它每次都必须从权重中重新计算，它还显示了一些严重的开销，可能是为了彻底的输入检查，不确定。无论如何，如果你把这个函数分解成小步骤（cdf、随机数生成、逆 cdf、值映射），你可以保持简单并预先计算更多的东西，节省一些时间。如果使用相同的种子重置 RNG（当然输入以相同的顺序处理），这两种方法都会导致相同的输出。

有了参考代码，我得到的时间和你一样。有了这两个变化，处理速度提高了大约 8 倍，还不错。

%%timeit -n 1 -r 5

sevents = events.sort_index()    # ensure that get_loc below will not return a mask (slow)
seiv = sevents.id.values
sewv = sevents.weight.values

def randomizer(t):
    s = sevents.index.get_loc(t[0])    # either a slice (because of sort) or a scalar
    v = seiv[s]

    if isinstance(s, slice):
        w = sewv[s]
        cw = w.cumsum()    # cumulative weight (i.e. cdf)
        cw /= cw[-1]
        return lambda: v[np.searchsorted(cw, np.random.rand() + 1e-35)]    # inverse cdf
    else:
        return lambda: v    # only one event with this time

# create a per-minute random function returning id according to weights
randomizers = sevents.index.unique().to_frame().apply(randomizer, axis='columns', raw=True)

# for each observation, find random generator and call it
selections = randomizers.loc[observations].apply(lambda f: f())

---

 1.67 s ± 12.4 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)