python - 为 pandas DataFrame 滚动 idxmin/max

Question

我相信以下函数是 pandas DataFrame 滚动 argmin/max 的有效解决方案:

import numpy as np

def data_frame_rolling_arg_func(df, window_size, func):
    ws = window_size
    wm1 = window_size - 1
    return (df.rolling(ws).apply(getattr(np, f'arg{func}'))[wm1:].astype(int) +
            np.array([np.arange(len(df) - wm1)]).T).applymap(
                lambda x: df.index[x]).combine_first(df.applymap(lambda x: np.NaN))

灵感来自a partial solution for rolling idxmax on pandas Series .

解释:

• 将 numpy argmin/max 函数应用于滚动窗口。
• 仅保留非NaN 值。
• 将值转换为 int。
• 将值重新对齐到原始行号。
• 使用 applymap 将行号替换为索引值。
• 与填充有 NaN 的原始 DataFrame 结合，以添加具有预期 NaN 值的第一行。

In [1]: index = map(chr, range(ord('a'), ord('a') + 10))

In [2]: df = pd.DataFrame((10 * np.random.randn(10, 3)).astype(int), index=index)

In [3]: df                                                                                                                                                                                                                                                                       
Out[3]: 
    0   1   2
a  -4  15   0
b   0  -6   4
c   7   8 -18
d  11  12 -16
e   6   3  -6
f  -1   4  -9
g   6 -10  -7
h   8  11 -25
i  -2 -10  -8
j   0  10  -7

In [4]: data_frame_rolling_arg_func(df, 3, 'max')                                                                                                                                                                                                                                
Out[4]: 
     0    1    2
a  NaN  NaN  NaN
b  NaN  NaN  NaN
c    c    a    b
d    d    d    b
e    d    d    e
f    d    d    e
g    e    f    e
h    h    h    g
i    h    h    g
j    h    h    j

In [5]: data_frame_rolling_arg_func(df, 3, 'min')                                                                                                                                                                                                                                
Out[5]: 
     0    1    2
a  NaN  NaN  NaN
b  NaN  NaN  NaN
c    a    b    c
d    b    b    c
e    e    e    c
f    f    e    d
g    f    g    f
h    f    g    h
i    i    g    h
j    i    i    h

我的问题是:

• 你能找出任何错误吗？
• 有更好的解决方案吗？即:性能更高和/或更优雅。

对于那里的 pandas 维护者:如果已经很棒的 pandas 库包含滚动 idxmax 和 idxmin 就好了。

Answer 1

我在对 OP 的评论中提到的 NaN 问题可以通过以下方式解决:

import numpy as np
import pandas as pd


def data_frame_rolling_idx_func(df, window_size, func):
    ws = window_size
    wm1 = window_size - 1
    return (df.rolling(ws, min_periods=0).apply(getattr(np, f'arg{func}'),
                                                raw=True)[wm1:].astype(int) +
            np.array([np.arange(len(df) - wm1)]).T).applymap(
                lambda x: df.index[x]).combine_first(df.applymap(lambda x: np.NaN))


def main():
    index = map(chr, range(ord('a'), ord('a') + 10))
    df = pd.DataFrame((10 * np.random.randn(10, 3)).astype(int), index=index)
    df[0][3:6] = np.NaN
    print(df)
    print(data_frame_rolling_arg_func(df, 3, 'min'))
    print(data_frame_rolling_arg_func(df, 3, 'max'))


if __name__ == "__main__":
    main()

结果:

$ python demo.py 
      0   1   2
a   3.0   0   7
b   1.0   3  11
c   1.0  15  -6
d   NaN   2 -16
e   NaN   0  24
f   NaN   0  14
g   2.0   0   4
h  -1.0 -11  16
i  17.0   0  -2
j   3.0  -5  -8
     0    1    2
a  NaN  NaN  NaN
b  NaN  NaN  NaN
c    b    a    c
d    d    d    d
e    d    e    d
f    d    e    d
g    e    e    g
h    f    h    g
i    h    h    i
j    h    h    j
     0    1    2
a  NaN  NaN  NaN
b  NaN  NaN  NaN
c    a    c    b
d    d    c    b
e    d    c    e
f    d    d    e
g    e    e    e
h    f    f    h
i    i    g    h
j    i    i    h

NaN 值的处理有点微妙。我希望我的滚动 idxmin/max 函数能够与常规 DataFrame 滚动 min/max 函数很好地协作。默认情况下，只要窗口输入显示 NaN 值，它们就会生成一个 NaN 值。默认情况下，滚动 apply 功能也是如此。但是对于 apply 函数来说，这是个问题，因为我无法将 NaN 值转换为索引。然而，这很遗憾，因为输出中的 NaN 值显示出来是因为它们可以在输入中找到，所以输入中的 NaN 值索引就是我想要的就像我要生成的滚动 idxmin/max 函数一样。幸运的是，如果我使用以下参数组合，这正是我将得到的结果:

• min_periods=0 用于 pandas rolling 功能。然后，apply 函数将有机会生成自己的值，而不管在输入窗口中找到了多少个 NaN 值。
• raw=True 用于apply 函数。此参数确保应用函数的输入作为 numpy 数组而不是 pandas 系列传递。 np.argmin/max 将返回第一个输入 NaN 值的索引，这正是我们想要的。应该注意的是，如果没有 raw=True，即在 pandas Series 的情况下，np.argmin/max 似乎忽略了 NaN 值，这不是我们想要的。 raw=True 的好处在于它也应该提高性能!稍后会详细介绍。