Python Pandas Group apply

user_15615416 0
Viewed 66

I have a function that I am using very frequently. I have a set of data that I want to group by certain columns and then apply a function to it to get a new dataframe out of it.

Lets say I have a dataset with three columns:

Channel | threshold | Counts

I want to group by channel number and make a Gaussian fit of each group. Then get a value for each group.

df_gauss  = grp.group_apply(df_new, ["Channel"] ,['mean', 'std',"A"],gauss_fit, ["thr","scaler"] ,[100., 3400., 1.] )

The implementation of the function looks like this and works fine, I was just wondering if there is a better way of doing this. In particular, is there any good way to avoid the "hack" with the dummyVariableName?

import numpy
from scipy.optimize import curve_fit
import pandas as pd

def gauss(x, *p):
    A, mu, sigma = p
    return A*numpy.exp(-(x-mu)**2/(2.*sigma**2))

def gauss_fit(x,y, p0):
    coeff, var_matrix = curve_fit(gauss,x, y, p0=p0)
    return coeff[1], coeff[2],coeff[0]

def group_apply(df, groupAxis, output_axis, func, input_axis, *args, **kwargs):
    def wrapped_fun(df, input_axis,*args, **kwargs):
        df_input = [df[x].to_numpy() for x in input_axis]
        return func(*df_input, *args, **kwargs)
    df1 = df.groupby(groupAxis).apply(wrapped_fun, input_axis, *args, **kwargs)  
    dummyVariableName = 'internal_dummy_name'
    df2 = df1.reset_index(name=dummyVariableName)
    df2[output_axis] = pd.DataFrame(df2[dummyVariableName].to_list(), columns=output_axis)
    df2 = df2.drop([dummyVariableName],axis=1)
    return df2
pandas numpy python
1 Answers

You can change wrapped_fun to return a Series instead of a tuple to avoid the trick with the internal_dummy_name. This way, you can also use add as_index=False to the groupby to avoid using any reset_index.

Code:

def group_apply(df, groupAxis, output_axis, func, input_axis, *args, **kwargs):
    def wrapped_fun(df):
        df_input = [df[x].to_numpy() for x in input_axis]
        return pd.Series(func(*df_input, *args, **kwargs), index=output_axis)
    return df.groupby(groupAxis, as_index=False).apply(wrapped_fun)

Note that there is no need to have input_axis, *args, and **kwargs as parameters for wrapped_fun since these are already "in scope".

user_75795470
Related