使用依赖于 colname 的 Pandas Rolling Apply 的自定义函数

Question

使用 Pandas 1.1.5，我有一个如下所示的测试 DataFrame：

import numpy as np
import pandas as pd
df = pd.DataFrame({'id': ['a0','a0','a0','a1','a1','a1','a2','a2'],
                   'a': [4,5,6,1,2,3,7,9],
                   'b': [3,4,5,3,2,4,1,3],
                   'c': [7,4,3,8,9,7,4,6],
                   'denom_a': [7,8,9,7,8,9,7,8],
                   'denom_b': [10,11,12,10,11,12,10,11]})

我想在滚动窗口上应用以下自定义聚合函数，其中函数的计算取决于列名：

def custom_func(s, df, colname):
  if 'a' in colname:
    denom = df.loc[s.index, "denom_a"]
    calc = s.sum() / np.max(denom)
  elif 'b' in colname:
    denom = df.loc[s.index, "denom_b"]
    calc = s.sum() / np.max(denom)
  else:
    calc = s.mean()
  return calc

df.groupby('id')\
  .rolling(2, 1)\
  .apply(lambda x: custom_func(x, df, x.name))

这会导致TypeError: argument of type 'NoneType' is not iterable，因为每列的窗口子集不保留原始df 列的名称。也就是说，作为参数传入的x.name实际上是在传递None，而不是原始列名的字符串。

是否有某种方法可以使这种方法发挥作用（例如，保留使用 apply 作用的列名并将其传递给函数）？或者有什么改变它的建议吗？我查阅了以下参考资料，让自定义函数在同一窗口计算中使用多个列，等等：

https://stackoverflow.com/a/57601839/6464695

Answer 1

如果有一个“更好”的解决方案，我不会感到惊讶，但我认为至少可以是一个“好的开始”（我对.rolling(...) 做的不多）。

对于这个解决方案，我做了两个关键假设：

1. 所有denom_<X> 都有一个对应的<X> 列。
2. 您对(<X>, denom_<X>) 对所做的一切都是一样的。 （这应该很容易根据需要进行自定义。）

话虽如此，我在函数内部而不是外部执行.rolling，部分原因是RollingGroupBy 上的.apply(...) 似乎只能按列工作，这在这里没有太大帮助（海事组织）。

def cust_fn(df: pd.DataFrame, rolling_args: Tuple) -> pd.Series:
    cols = df.columns
    denom_cols = ["id"]  # the whole dataframe is passed, so place identifiers / uncomputable variables here.

    for denom_col in cols[cols.str.startswith("denom_")]:
        denom_cols += [denom_col, denom_col.replace("denom_", "")]
        col = denom_cols[-1]  # sugar
        df[f"calc_{col}"] = df[col].rolling(*rolling_args).sum() / df[denom_col].max()

    for col in cols[~cols.isin(denom_cols)]:
        print(col, df[col])
        df[f"calc_{col}"] = df[col].rolling(*rolling_args).mean()
    
    return df

那么你运行它的方式如下（你会得到相应的输出）：

>>> df.groupby("id").apply(cust_fn, rolling_args=(2, 1))
   id  a  b  c  denom_a  denom_b    calc_a    calc_b  calc_c
0  a0  4  3  7        7       10  0.444444  0.250000     7.0
1  a0  5  4  4        8       11  1.000000  0.583333     5.5
2  a0  6  5  3        9       12  1.222222  0.750000     3.5
3  a1  1  3  8        7       10  0.111111  0.250000     8.0
4  a1  2  2  9        8       11  0.333333  0.416667     8.5
5  a1  3  4  7        9       12  0.555556  0.500000     8.0
6  a2  7  1  4        7       10  0.875000  0.090909     4.0
7  a2  9  3  6        8       11  2.000000  0.363636     5.0

---

如果您需要动态说明存在哪些非数字/可计算列，那么将cust_fn 定义如下：

def cust_fn(df: pd.DataFrame, rolling_args: Tuple, index_cols: List = []) -> pd.Series:
    cols = df.columns
    denon_cols = index_cols

    # ... the rest is unchanged

然后您将调整您对cust_fn 的调用，如下所示：

>>> df.groupby("id").apply(cust_fn, rolling_args=(2, 1), index_cols=["id"])

---

当然，如果您在调整它以适应您的用途时遇到问题，请对此发表评论。 ?