在Pandas groupby().agg() 中使用闭包计算加权平均值

本文详细介绍了在pandas `groupby().agg()` 操作中，如何解决在聚合函数内部访问外部dataframe上下文的`nameerror`问题。通过引入python闭包（closure）机制，我们能够优雅地将外部dataframe作为权重参数传递给自定义的加权平均函数，从而实现复杂的数据聚合逻辑，确保代码的正确性和可维护性。

理解groupby().agg()中的上下文限制

在使用Pandas进行数据分析时，groupby().agg()是一个强大且常用的操作，它允许我们对数据进行分组，并对每个组应用一个或多个聚合函数。然而，当自定义聚合函数需要访问分组之外的原始DataFrame中的其他列作为辅助信息（例如计算加权平均值时需要权重列）时，常常会遇到NameError。

例如，在计算other_col的加权平均值时，权重信息amount位于原始DataFrame的另一列。如果直接在聚合函数weighted_mean中尝试访问全局变量df1或期望df1能自动传递，就会出现NameError: name 'df1' is not defined。这是因为agg方法在执行时，会将其作用的Series（例如other_col的每个组）传递给聚合函数，而聚合函数本身并不直接拥有原始DataFrame的完整上下文。它只接收到当前分组的Series数据x。

解决方案：利用Python闭包

Python中的闭包（closure）提供了一种优雅的解决方案。闭包是指一个函数记住其创建时的环境，即使该环境已经不存在，它仍然可以访问该环境中的非局部变量。我们可以利用这一特性，创建一个外部函数来“捕获”原始DataFrame，并返回一个内部函数作为实际的聚合函数。

闭包工作原理

1. 外部函数接收上下文： 定义一个外部函数（例如weighted_mean_factory），它接收原始DataFrame作为参数。
2. 内部函数执行聚合： 在外部函数内部定义一个内部函数（例如inner_weighted_mean），这个内部函数将作为agg方法的实际聚合函数，接收分组后的Series x。
3. 内部函数访问外部变量： 内部函数可以访问外部函数作用域中的DataFrame（即闭包捕获的变量）。
4. 外部函数返回内部函数： 外部函数返回这个内部函数。

当外部函数被调用时，它会创建一个包含原始DataFrame上下文的内部函数实例。这个内部函数随后被传递给agg方法，当agg调用它时，它依然能够访问到被捕获的DataFrame，从而获取权重信息。

示例代码与实现

下面我们将通过一个具体的例子来演示如何使用闭包解决上述问题。

假设我们有一个DataFrame df2，包含id、amount和other_col三列。我们的目标是按id分组，计算amount的总和是否大于100（xx列），以及other_col的加权平均值（权重为amount列），并判断其是否大于0.5（yy列）。

算家云

高效、便捷的人工智能算力服务平台

37

查看详情

import pandas as pd
import numpy as np

# 原始DataFrame
df2 = pd.DataFrame({
    'id': [1, 1, 2, 2, 3],
    'amount': [10, 200, 1, 10, 150],
    'other_col': [0.1, 0.6, 0.7, 0.2, 0.4]
})

def weighted_mean_factory(df_context):
    """
    这是一个闭包工厂函数，它接收原始DataFrame作为上下文，
    并返回一个可以用于groupby().agg()的聚合函数。
    """
    def inner_weighted_mean(x):
        """
        实际执行加权平均计算的内部函数。
        x 是当前分组的 Series (例如 'other_col' 的一个分组)。
        """
        try:
            # 使用 df_context 访问原始DataFrame，通过 x.index 获取对应行的权重
            weighted_avg = np.average(x, weights=df_context.loc[x.index, 'amount'])
            return weighted_avg > 0.5
        except ZeroDivisionError:
            # 处理权重总和为零的情况，避免除以零错误
            return 0
    return inner_weighted_mean

def some_function(df_input):
    """
    主函数，用于执行分组聚合操作。
    """
    # 创建一个带有 df_input 上下文的加权平均函数实例
    weighted_mean_for_current_df = weighted_mean_factory(df_input)

    # 执行分组聚合
    df_result = df_input.groupby('id').agg(
        xx=('amount', lambda x: x.sum() > 100),  # 检查 amount 总和是否大于 100
        yy=('other_col', weighted_mean_for_current_df) # 使用闭包返回的函数
    ).reset_index()
    return df_result

# 调用主函数并打印结果
df_processed = some_function(df_input=df2)
print(df_processed)

代码解析

1. weighted_mean_factory(df_context):

   • 这个外部函数接收一个参数df_context，它就是我们希望闭包捕获的原始DataFrame。
   • 它定义了一个内部函数inner_weighted_mean(x)。
   • 最后，它返回inner_weighted_mean。

2. inner_weighted_mean(x):

   • 这个函数是真正被groupby().agg()调用的函数。x是当前分组的other_col Series。
   • 关键在于df_context.loc[x.index, 'amount']。x.index提供了当前分组元素在原始DataFrame中的索引。通过这些索引，我们可以精确地从df_context中取出对应的amount值作为权重。
   • np.average()函数用于计算加权平均值。
   • try-except ZeroDivisionError块用于处理当weights总和为零时可能发生的除以零错误，提供更健壮的代码。

3. some_function(df_input):

   • 这个函数负责协调整个流程。
   • weighted_mean_for_current_df = weighted_mean_factory(df_input)：在这里，我们调用weighted_mean_factory，并将df_input（即df2）传递给它。weighted_mean_factory返回一个inner_weighted_mean的实例，这个实例已经“记住”了df_input。
   • df_input.groupby('id').agg(...)：在agg方法中，我们将这个weighted_mean_for_current_df函数作为yy列的聚合函数。当agg调用它时，它能够正确访问到df_input上下文。

预期输出

运行上述代码，将得到以下结果：

   id     xx     yy
0   1   True   True
1   2  False  False
3   3   True  False

注意事项与总结

• 性能考量： 对于非常大的DataFrame，在agg函数内部频繁地使用df_context.loc[x.index, 'column_name']可能会有一定的性能开销，因为loc操作在每次分组时都会执行。在某些极端性能敏感的场景下，可能需要考虑其他优化策略，例如提前计算并合并权重列。然而，对于大多数常见用例，闭包的这种方式已经足够高效且代码可读性强。
• 代码可读性： 使用闭包模式能够清晰地分离聚合逻辑与上下文数据，提高了代码的可读性和模块化。
• 适用性： 闭包不仅适用于加权平均，还可以用于任何需要在groupby().agg()内部访问原始DataFrame其他列信息的复杂聚合逻辑。

通过理解和应用Python闭包，我们可以有效地解决Pandas groupby().agg()中聚合函数上下文缺失的问题，从而实现更灵活、更强大的数据分析功能。

以上就是在Pandas groupby().agg() 中使用闭包计算加权平均值的详细内容，更多请关注php中文网其它相关文章！