Pandas groupby weighted average-CodePudding

I have a DataFrame on which I would like to group by date ("count_date" in column in my df) and apply a weighted average on the average speed ("average_speed") weighted by the count ("count").

I'm trying this:

df_byday = df_merged.groupby("count_date").apply(lambda x: np.average(x['average_speed'], weights=x['count']))

but it returns a ZeroDivisionError as there are some rows where "count" and "average_speed" are equal to 0.

While searching in the Numpy documentation I could see that np.ma.average() could help solve the issue. But when I'm trying to apply it, I have the following error:

/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/numpy/ma/extras.py:649: RuntimeWarning: invalid value encountered in double_scalars
  avg = np.multiply(a, wgt,

Could anyone help and let me know how to deal with ZeroDivisionError in this case?

Thank you very much!

CodePudding user response：

You can catch ZeroDivisionError exception

def func(x):
    try: 
        return np.average(x['average_speed'], weights=x['count'])
    except ZeroDivisionError:
        return 0

df_byday = df_merged.groupby("count_date").apply(func)