Средневзвешенные панды

Question

Я вычисляю взвешенное среднее для многих столбцов, используя панд. В некоторых случаях вес может равняться нулю, поэтому я использую np.ma.average:

import pandas as pd
import numpy as np

df = pd.DataFrame.from_dict(dict([('ID', [1, 1, 1]),('HeightA', [1, 2, 3]), ('WeightA', [0, 0, 0]),('HeightB', [2, 4, 6]), ('WeightB', [1, 2, 4])]))

>>df
   ID  HeightA  WeightA  HeightB  WeightB
0   1        1        0        2        1
1   1        2        0        4        2
2   1        3        0        6        4


wmA = lambda x: np.ma.average(x, weights=df.loc[x.index, "WeightA"])
wmB = lambda x: np.ma.average(x, weights=df.loc[x.index, "WeightB"])
f = {'HeightA':wmA,'HeightB':wmB}
df2 = df.groupby(['ID'])['HeightA','HeightB'].agg(f)

Это работает, но у меня есть много столбцов высоты и веса, поэтому я не хочу писать лямбда-функцию для каждого, поэтому я пытаюсь:

def givewm(data,weightcolumn):
    return np.ma.average(data, weights=data.loc[data.index, weightcolumn])

f = {'HeightA':givewm(df,'WeightA'),'HeightB':givewm(df,'WeightB')}
df2 = df.groupby(['ID'])['HeightA','HeightB'].agg(f)

Что выдает ошибку: builtins.TypeError: Ось должна быть указана, когда формы и веса различаются.

Как мне написать функцию, которая возвращает взвешенное среднее с именем столбца веса в качестве входных данных?

Answer 1

Используйте двойные вложенные функции, решение от github :

df = (pd.DataFrame.from_dict(dict([('ID', [1, 1, 1]),
                                  ('HeightA', [1, 2, 3]), 
                                  ('WeightA', [10, 20, 30]),
                                  ('HeightB', [2, 4, 6]), 
                                  ('WeightB', [1, 2, 4])])))


print (df)
   ID  HeightA  WeightA  HeightB  WeightB
0   1        1       10        2        1
1   1        2       20        4        2
2   1        3       30        6        4


def givewm(weightcolumn):
    def f1(x):
        return np.ma.average(x, weights=df.loc[x.index, weightcolumn])
    return f1

f = {'HeightA':givewm('WeightA'),'HeightB':givewm('WeightB')}
df2 = df.groupby('ID').agg(f)
print (df2)
     HeightA   HeightB
ID                    
1   2.333333  4.857143

Проверить решение:

wmA = lambda x: np.ma.average(x, weights=df.loc[x.index, "WeightA"])
wmB = lambda x: np.ma.average(x, weights=df.loc[x.index, "WeightB"])
f = {'HeightA':wmA,'HeightB':wmB}

df2 = df.groupby(['ID'])['HeightA','HeightB'].agg(f)
print (df2)
     HeightA   HeightB
ID                    
1   2.333333  4.857143