加快 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 上的下限来设置相应的行元素。

以 pandas 表 D 为例

   C1  C2  C3
0  2   5   5
1  10  12   2
2   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'] = 16。我们看到第一个这样的 i 是 1 并且因为这恰好满足第一个条件，所以我们将这一行的新列设置为 0。

不幸的是，上述解决方案效率很低。我尝试通过向后计算值并尝试缓存结果来优化它（有时可以从“过去”值推断出 C4 的值），但不幸的是它并没有明显好转。

根据我的经验，获得最大性能的最佳方法是尝试在 pandas 框架内尽可能多地表达。

有没有什么有意义的方法可以优化上面的代码？

编辑。使用已接受答案的代码并替换以下函数可获得最佳结果。

@njit
def 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

您可能会发现 pandas 函数的用途：IntervalIndex，参考如下：

data.index = pd.IntervalIndex.from_tuples(data['C1'], data['C3'])
I = data.index

你的函数可能被重写为：

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

   for c in coeff:

        i = I.map(c)

        low_bound = i - rate_low*c
        high_bound = i + rate_high*c

        for j in range(k+1, len(data)):
            if data.at[ I.get_loc(j), 'C2'] < low_bound:
                data.at[i, bot_col_name] = 0 
                break

            if data.at[ I.get_loc(j), 'C3'] >= high_bound:
                data.at[i, bot_col_name] = 1 
                break

如果这有帮助，请告诉我！

Answer 2

您确定算法完全按照您的想法运行吗？特别是内部循环，您在其中查找具有 C2 = 上限的下一行（例如，如果同一行同时满足两个条件会发生什么情况？）。无论如何，假设它是，通常加快 Pandas 中任何算法的关键，正如您所暗示的，将所有内容放入 DataFrame，然后使用基于列的操作。我建议如下：

# Setup DataFrame
df = pd.DataFrame({
    "C1": [2, 10, 8, 30],  
    "C2": [5, 12, 3, 25],
    "C3": [5, 2, 17, 3]
})

coeffs = [3, 3, 5, 7]

df['coeff'] = coeffs
df['C1_next'] = df['C1'].shift(-1)


# Fixed Values
rate_low = 2
rate_high = 2 


# Helper Functions
def test_bounds(row, lower, rate):

    if lower:
        bound = row['C1_next'] - rate * row['coeff']
    elif not lower:
        bound = row['C1_next'] + rate * row['coeff']

    return bound


def earliest_bound(x, lower):
    rows_to_search = df[df.index > x.name]

    if lower:
        slice = rows_to_search[rows_to_search['C2'] < x['lower_bound']]
    elif not lower:
        slice = rows_to_search[rows_to_search['C3'] >= x['higher_bound']]

    if len(slice) > 0:
        value = slice.index[0]
    else:
        value = np.NaN

    return value



df['lower_bound'] = df.apply(test_bounds, args=(True, rate_low), axis=1)
df['higher_bound'] = df.apply(test_bounds, args=(False, rate_high), axis=1)

df["earliest_lower_bound_row"] = df.apply(earliest_bound, args=(True, ), axis=1)
df["earliest_higher_bound_row"] = df.apply(earliest_bound, args=(False, ), axis=1)

在这种情况下，early_lower_bound_row 和 early_higher_bound_row 中返回的值将是满足该条件的第一行。当然，如果您只想要一个实际的 0/1 值，那么可以直接使用这两列中的信息来创建。

Answer 3

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

Answer 4

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

将代码编译为 c / Fortran 是 numpy / pandas 内部函数如此快速的原因。更多信息请访问pandas documentation。

让我们首先创建示例：

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 没有完成。