Pandas groupby agg apply a function with multiple parameters

Question

How can a function be applied on a pandas groupby that requires parameters from multiple columns of the groupby dataframe and returns two scaler values.

Below is the repeatable example. The last line gets the f_value

import pandas as pd
import numpy as np
from statsmodels.formula.api import ols
import plotly.express as px

n=100
df = pd.DataFrame({
    'c': np.random.choice(['CATS', 'DOGS'], n),
    'x': np.random.choice(list('ABCDE'), n),
    'y': np.random.normal(5, 1, n)
})

signal = np.where(df['c'].eq('CATS') & df['x'].eq('A'), 1.1, 0)
df['y'] = df['y'] + signal

def get_ols_fp(df, x, y):
    formula = y + '~' + x
    model = ols(formula, df).fit()
    f_value = model.fvalue
    p_value = model.f_pvalue
    return (f_value, p_value)

# getting f_value and p_value works with a single series.
get_ols_fp(df[df['c'].eq('CATS')], 'x', 'y')

This above code works and fetches the f_value and the p_value. However, the following does not work.

# how could we run the get_ols with a groupby().agg() 
df.groupby('c').agg(get_ols_fp('x', 'y'))

The desired output would be a dataframe one row per level of the 'c' variable ('CATTS' and 'DOGS') in this case and one column for the p_value, and another for the f_value.

user_8134197 · Answer

This is working :

def get_ols_fp(df, x=None, y=None):
    formula = y + '~' + x 
    model = ols(formula, df).fit()
    f_value = model.fvalue
    p_value = model.f_pvalue
    return pd.Series([f_value, p_value], index=['f_value', 'p_value'])

df.groupby('c').apply(get_ols_fp, x='x', y = 'y')

user_2411555 · Answer

I'd do it a little different. I don't know if it's the easiest way, but it works.

Example:

import pandas as pd
import numpy as np
from statsmodels.formula.api import ols

n=100
df = pd.DataFrame({
    'c': np.random.choice(['CATS', 'DOGS'], n),
    'x': np.random.choice(list('ABCDE'), n),
    'y': np.random.normal(5, 1, n)
})

signal = np.where(df['c'].eq('CATS') & df['x'].eq('A'), 1.1, 0)
df['y'] = df['y'] + signal

def get_ols_fp(df, x, y):
    formula = y + '~' + x
    model = ols(formula, df).fit()
    f_value = model.fvalue
    p_value = model.f_pvalue
    return (f_value, p_value)

# getting f_value and p_value works with a single series.
# get_ols_fp(df[df['c'].eq('CATS')], 'x', 'y')

df_result = pd.DataFrame([], columns = ["c", "f_value", "p_value"])
for c, dd in df.groupby(['c']):
    v = get_ols_fp(dd, 'x', 'y')
    df_result.loc[len(df_result)] = [c, *v]

df_result

Pandas groupby agg apply a function with multiple parameters

2 Answers