Объединить / объединить / объединить функции в Python

Question

Я создал функцию для создания тепловых карт, и данные, которые я передаю, могут быть сгруппированы по разным столбцам и агрегированы с разными функциями (среднее, сумма, количество). Вплоть до того, что у меня была конфигурация if-elif, где условие должно соответствовать строке требуемой функции (например, если 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]

Дело в том, что значения сами по себе не функционируют, и этот подход не работает. Так что я хотел бы вызвать функцию, которая доступна для результата группирования, для агрегирования, установить функцию, которую я хочу передать, или вызвать словарь. Если это конкретное имя, пожалуйста, дайте мне знать, чтобы я мог переименовать вопрос более точно.

Answer 1

Определите вашу функцию следующим образом:

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

Преимущество заключается в вызове функции модуля с использованием ее имени (строки). Причина, по которой ваш диктант не может работать, состоит в том, что вы не назначили правильную функцию, которая должна быть ограничена внутри модуля и не должна вызываться.