如何在Pandas Groupby中使用带多参数的自定义聚合函数？

问题背景 #

我已经能在Pandas的groupby中成功应用仅接收分组值的自定义函数，但现在需要使用一个带额外参数的函数。

我的DataFrame结构如下：

a     b     c      value
a1    b1    c1      v1
a2    b2    c2      v2
a3    b3    c3      v3

单参数版本的可运行代码：

import numpy as np
import pandas as pd

def cpk(a):
    arr = np.asarray(a)
    arr = arr.ravel()
    sigma = np.std(arr)
    m = np.mean(arr)

    Cpu = float(150 - m) / (3*sigma)
    Cpl = float(m - 50) / (3*sigma)
    Cpk = np.min([Cpu, Cpl])
    return Cpk

# 调用方式
df_cpk = df_result.groupby(['a','b','c'])['value'].agg(cpk).reset_index()

现在需要应用带多参数的cpk2函数，该函数除了分组值外，还需要传入lsl（下限规格）和usl（上限规格）：

def cpk2(a,lsl,usl):
    arr = np.asarray(a)
    arr = arr.ravel()
    sigma = np.std(arr)
    m = np.mean(arr)

    Cpu = float(usl - m) / (3*sigma)
    Cpl = float(m - lsl) / (3*sigma)
    Cpk = np.min([Cpu, Cpl])
    return Cpk

请问有什么简便的实现方法？

解决方案 #

方法1：使用lambda表达式传递参数 #

直接用lambda把分组值和额外参数一起传给cpk2，写法直观：

# 假设lsl=50，usl=150
df_cpk = df_result.groupby(['a','b','c'])['value'].agg(lambda x: cpk2(x, lsl=50, usl=150)).reset_index()

方法2：使用functools.partial固定参数 #

如果需要重复使用相同的lsl和usl，可以用partial预先绑定参数，简化后续调用：

from functools import partial

# 固定lsl和usl的值
cpk_partial = partial(cpk2, lsl=50, usl=150)
df_cpk = df_result.groupby(['a','b','c'])['value'].agg(cpk_partial).reset_index()

方法3：利用Pandas的agg直接传关键字参数（Pandas 0.25+支持） #

从Pandas 0.25版本开始，agg支持直接传递函数的关键字参数，写法更简洁：

df_cpk = df_result.groupby(['a','b','c'])['value'].agg(cpk2, lsl=50, usl=150).reset_index()

额外注意事项 #

如果分组内的value全部相同，会导致sigma=0，触发除以0的错误，可以在函数中加入判断逻辑避免：

def cpk2(a,lsl,usl):
    arr = np.asarray(a)
    arr = arr.ravel()
    sigma = np.std(arr)
    if sigma == 0:
        return 0  # 可根据业务需求调整返回值
    m = np.mean(arr)

    Cpu = float(usl - m) / (3*sigma)
    Cpl = float(m - lsl) / (3*sigma)
    Cpk = np.min([Cpu, Cpl])
    return Cpk

内容的提问来源于stack exchange，提问作者hjsg1010