groupby weighted average and sum in pandas dataframe

The solution that uses a dict of aggregation functions will be deprecated in a future version of pandas (version 0.22):

FutureWarning: using a dict with renaming is deprecated and will be removed in a future 
version return super(DataFrameGroupBy, self).aggregate(arg, *args, **kwargs)

Use a groupby apply and return a Series to rename columns as discussed in: Rename result columns from Pandas aggregation ("FutureWarning: using a dict with renaming is deprecated")

def my_agg(x):
    names = {'weighted_ave_price': (x['adjusted_lots'] * x['price']).sum()/x['adjusted_lots'].sum()}
    return pd.Series(names, index=['weighted_ave_price'])

produces the same result:

>df.groupby(["contract", "month", "year", "buys"]).apply(my_agg)

                          weighted_ave_price
contract month year buys                    
C        Z     5    Sell          424.828947
CC       U     5    Buy          3328.000000
SB       V     5    Buy            11.637500
W        Z     5    Sell          554.850000

Wouldn't it be a lot more simpler to do this.

  1. Multiply (adjusted_lots * price_weighted_mean) into a new column "X"
  2. Use groupby().sum() for columns "X" and "adjusted_lots" to get grouped df df_grouped
  3. Compute weighted average on the df_grouped as df_grouped['X']/df_grouped['adjusted_lots']

Doing weighted average by groupby(...).apply(...) can be very slow (100x from the following). See my answer (and others) on this thread.

def weighted_average(df,data_col,weight_col,by_col):
    df['_data_times_weight'] = df[data_col]*df[weight_col]
    df['_weight_where_notnull'] = df[weight_col]*pd.notnull(df[data_col])
    g = df.groupby(by_col)
    result = g['_data_times_weight'].sum() / g['_weight_where_notnull'].sum()
    del df['_data_times_weight'], df['_weight_where_notnull']
    return result

EDIT: update aggregation so it works with recent version of pandas

To pass multiple functions to a groupby object, you need to pass a tuples with the aggregation functions and the column to which the function applies:

# Define a lambda function to compute the weighted mean:
wm = lambda x: np.average(x, weights=df.loc[x.index, "adjusted_lots"])

# Define a dictionary with the functions to apply for a given column:
# the following is deprecated since pandas 0.20:
# f = {'adjusted_lots': ['sum'], 'price': {'weighted_mean' : wm} }
# df.groupby(["contract", "month", "year", "buys"]).agg(f)

# Groupby and aggregate with namedAgg [1]:
df.groupby(["contract", "month", "year", "buys"]).agg(adjusted_lots=("adjusted_lots", "sum"),  
                                                      price_weighted_mean=("price", wm))

                          adjusted_lots  price_weighted_mean
contract month year buys                                    
C        Z     5    Sell            -19           424.828947
CC       U     5    Buy               5          3328.000000
SB       V     5    Buy              12            11.637500
W        Z     5    Sell             -5           554.850000

You can see more here:

  • http://pandas.pydata.org/pandas-docs/stable/groupby.html#applying-multiple-functions-at-once

and in a similar question here:

  • Apply multiple functions to multiple groupby columns

Hope this helps

[1] : https://pandas.pydata.org/pandas-docs/stable/whatsnew/v0.25.0.html#groupby-aggregation-with-relabeling

Tags:

Python

Pandas

R