加快 pandas 迭代，以测试后续元素的条件

Question

给定一个具有三列 (C1,C2,C3) 和一系列等长正数 (coeff) 的 pandas 数据帧，我正在计算第四列 C4，如下所示def event(data, coeff, rate_low=2, rate_high=2):    bot_col_name = 'C4'    data[bot_col_name] = -1    I = data.index     for k in range(len(I)-1):        i = I[k]        next_val = data.at[ I[k+1], 'C1']        c = coeff.at[i]        low_bound = next_val - rate_low*c        high_bound = next_val + rate_high*c        for j in range(k+1, len(data)):            if data.at[ I[j], 'C2'] < low_bound:                data.at[i, bot_col_name] = 0                 break            if data.at[ I[j], 'C3'] >= high_bound:                data.at[i, bot_col_name] = 1                 break    return data换句话说，给定一行，我们计算某个上限和下限，然后根据我们首先达到 C2 下的上限还是 C3 上的下限来设置相应的行元素。作为一个例子，考虑熊猫表 D   C1  C2  C30  2   5   51  10  12   22   8   3  17 3  30  25   3现在如果 coeff = [3,3,5,7] 那么在计算第一行的值时，low_bound 为 10-2*3=4，high_bound 为 10+2*3=16。我们现在必须找到最小索引 i>0 使得 D.loc[i, 'C2'] < 4 或 D.loc[i,'C3'] >= 16。我们看到第一个这样的 i 是 1 并且因为这恰好满足第一个条件，所以我们将这一行的新列设置为 0。不幸的是，上述解决方案效率很低。我已经尝试通过向后计算值并尝试缓存结果来优化它（有时可以从“过去”值推断出 C4 的值），但不幸的是它并没有明显更好。根据我的经验，获得最大性能的最佳方法是尝试在 pandas 框架内尽可能多地表达。是否有任何有意义的方式可以优化上述代码？编辑。使用已接受答案的代码并替换以下函数可获得最佳结果。@njitdef get_c4(low_bound, high_bound, c2, c3):    r1 = np.argwhere( c2 < low_bound )    r2 = np.argwhere( c3 >= high_bound )    if len(r1) == 0 and len(r2) == 0:        return -1    elif len(r1) == 0:        return 1    elif len(r2) == 0:        return 0    return int (r1[0] > r2[0])

Answer 1

如果你真的需要一个快速的解决方案，你应该使用numba。numba的替代方案是cython。两者都编译你的 python 代码c以使其更快，但我认为numba更简单，它们或多或少具有相同的性能。

将代码编译到c/Fortran是numpy / pandas内部函数如此快速的原因。更多信息请参阅pandas 文档。

让我们首先创建示例：

import numpy as np

import pandas as pd

from numba import njit

df = pd.DataFrame({

    'C1': [2, 10, 8, 30],

    'C2': [5, 12, 3, 25],

    'C3': [5, 2, 17, 3]

})

coeff = pd.Series([3, 3, 5, 7])

然后通过转换为numba我们得到答案的代码：

@njit

def event_v(data, coeff, rate_low=2, rate_high=2):

    out = -np.ones(len(data), dtype=np.int8)

    for k in range(len(data) - 1):

        next_val = data[k + 1, 0]

        c = coeff[k]

        low_bound = next_val - rate_low * c

        high_bound = next_val + rate_high * c

        for j in range(k + 1, len(data)):

            if data[j, 1] < low_bound:

                out[k] = 0

                break

            if data[j, 2] >= high_bound:

                out[k] = 1

                break

    return out

df["C4"] = event_v(df.values, coeff.values)

测试 10 000 行：

n = 10_000

df = pd.DataFrame(np.random.randint(30, size=[n, 3]), columns=["C1", "C2", "C3"])

coeff = pd.Series(np.random.randint(10, size=n))

%timeit event_v(df.values, coeff.values)

3.39 ms ± 1.13 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit event(df, coeff) # Code from the question

28.4 s ± 1.02 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

它快了大约8500 倍

测试 1 000 000 行：

n = 1_000_000

df = pd.DataFrame(np.random.randint(30, size=[n, 3]), columns=["C1", "C2", "C3"])

coeff = pd.Series(np.random.randint(10, size=n))

%timeit event_v(df.values, coeff.values)

27.6 s ± 1.16 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

我尝试使用问题的代码运行它，但超过 2 小时%timeit后没有完成。

Answer 2

next_value、low_bound 和 high_bound 可以很容易地向量化，并且它们的计算速度非常快。第二部分不容易矢量化，因为它可能需要扫描整个数组的每一行。通过在 numpy 数组中进行比较，可以获得对您的实现的轻微改进（对于较大的 n 变得更加重要）。

def get_c4(low_bound, high_bound, c2, c3):

    for idx in range(len(c2)):

        if c2[idx] < low_bound:

            return 0

        if c3[idx] >= high_bound:

            return 1

    return -1

def event_new(data: pd.DataFrame, coeff, rate_low=2, rate_high=2):

    data['next_val'] = data['C1'].shift(periods=-1).ffill().astype('int')

    data['low_bound'] = (data['next_val'] - rate_low * coeff).astype('int')

    data['high_bound'] = (data['next_val'] + rate_high * coeff).astype('int')

    c2 = data['C2'].to_numpy()

    c3 = data['C3'].to_numpy()

    data['C4'] = data.apply(lambda x: get_c4(x.low_bound, x.high_bound, c2[data.index.get_loc(x) + 1:], c3[data.index.get_loc(x) + 1:]), axis=1)

    data.drop(columns=['next_val', 'low_bound', 'high_bound'])

    return data

基准代码：

for n in [1e2, 1e3, 1e4, 1e5, 1e6]:

    n = int(n)

    df = pd.DataFrame({'C1': random_list(n=n), 'C2': random_list(n=n), 'C3': random_list(n=n)})

    coeff = pd.Series(random_list(start=2, stop=7, n=n))

    print(f"n={n}:")

    print(f"Time org: {timeit.timeit(lambda: event(df.copy(), coeff), number=1):.3f} seconds")

    print(f"Time new: {timeit.timeit(lambda: event_new(df.copy(), coeff), number=1):.3f} seconds")

输出：

n=100:

Time org: 0.007 seconds

Time new: 0.012 seconds

n=1000:

Time org: 0.070 seconds

Time new: 0.048 seconds

n=10000:

Time org: 0.854 seconds

Time new: 0.493 seconds

n=100000:

Time org: 7.565 seconds

Time new: 4.456 seconds

n=1000000:

Time org: 216.408 seconds

Time new: 45.199 seconds