在python中连接/连接/聚合函数

Question

我创建了一个生成热图的函数，我输入的数据可以按不同的列进行分组，并使用不同的函数（mean、sum、count）进行聚合。 据我所知，我有一个if-elif配置，其中的条件是匹配所需函数的字符串（例如if-fun==“mean”）。你知道吗

我所拥有的：

df_ex = pd.DataFrame({"id_place":np.random.randint(20, size=133),
                    "id_center": np.random.randint(35, size=133),
                    "flow": np.random.randint(500, size=133),
                    "temperature": np.random.randint(35, size=133),
                    })
def prep_df(df,columns,variable, fun):

   columns= columns
   grouped_data = df.groupby(columns)[variable]

   if fun== 'mean':
       grouped_data = grouped_data.mean()
   elif fun == 'count':
       grouped_data = grouped_data.count()
   elif fun == 'sum':
       grouped_data = grouped_data.sum()

   return grouped_data

prep_df(df_ex,["id_place","id_center"], "flow", "mean")

为了使代码更简洁，我希望像字典这样的东西只有一行，而不是if elif所需的所有行，然后我可以通过将函数参数作为键来调用感兴趣的函数。 比如：

def prep_df(df,columns,variable, fun):
    dict_fun = {'mean': mean(),
                'count': count(),
                'sum': sum() 
                }
    grouped_data = df.groupby(columns)[variable].dict_fun[fun]

事实上，价值观本身并不起作用，这种方法失败了。 所以我想调用groupby结果中可用的函数进行聚合，设置传递变量或调用字典的函数。 这是一个具体的名字，请让我知道，这样我可以更名的问题更准确。你知道吗

Answer 1

定义函数如下：

def prep_df2(df, columns, variable, fun):
    assert fun in {'mean', 'count', 'sum'}
    return getattr(df.groupby(columns)[variable], fun)()

它利用了使用模块名（字符串）调用模块函数的优势。 dict不能工作的原因是您没有分配正确的函数，该函数应该限定在模块内部，不应该被调用。你知道吗