极坐标组当权重总和为零时使用等加权平均值

Question

我想在 groupby 操作期间计算极坐标数据帧的加权平均值。如果 groupby 内的权重总和等于零，我想返回同等加权的平均值。此外，我想忽略要平均的列或权重列中的空值。

记号：

df (polars.DataFrame): input dataframe
gr (list[str]): columns to groupby
vr (list[str]): columns to take average of
wt (str): column to use as the weights

计算我的等加权平均值

res = df.group_by(gr, maintain_order = True).agg(polars.col(vr).mean())

默认情况下，均值函数会忽略“vr”列中的空值，这是所需的。

计算我的加权平均值

mask_wt = polars.col(wt) * polars.col(vr).is_not_null()
wavg = (polars.col(vr) * polars.col(wt)).sum() / mask_wt.sum()
res = df.group_by(gr, maintain_order = True).agg(wavg)

屏蔽权重允许我忽略“vr”列和“wt”列中的空值。

例如，

import polars

df = polars.DataFrame({'id': ['x', 'x', 'y', 'y'], 
                       'dt': ['j', 'f', 'j', 'f'],
                       'a': [None, 2, 3, 4], 
                       'b': [5, 6, 7, 8], 
                       'wt': [1, 0, 0.25, 0.25]})

gr = ['id']
vr = ['a', 'b']
wt = 'wt'

>>> df
┌─────┬─────┬──────┬─────┬──────┐
│ id  ┆ dt  ┆ a    ┆ b   ┆ wt   │
│ --- ┆ --- ┆ ---  ┆ --- ┆ ---  │
│ str ┆ str ┆ i64  ┆ i64 ┆ f64  │
╞═════╪═════╪══════╪═════╪══════╡
│ x   ┆ j   ┆ null ┆ 5   ┆ 1.0  │
│ x   ┆ f   ┆ 2    ┆ 6   ┆ 0.0  │
│ y   ┆ j   ┆ 3    ┆ 7   ┆ 0.25 │
│ y   ┆ f   ┆ 4    ┆ 8   ┆ 0.25 │
└─────┴─────┴──────┴─────┴──────┘

>>> df.group_by(gr, maintain_order = True).agg(polars.col(vr).mean())
┌─────┬─────┬─────┐
│ id  ┆ a   ┆ b   │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═════╪═════╪═════╡
│ x   ┆ 2.0 ┆ 5.5 │
│ y   ┆ 3.5 ┆ 7.5 │
└─────┴─────┴─────┘

>>> mask_wt = polars.col(wt) * polars.col(vr).is_not_null()
>>> wavg = (polars.col(vr) * polars.col(wt)).sum() / mask_wt.sum()
>>> res = df.group_by(gr, maintain_order = True).agg(wavg)
┌─────┬─────┬─────┐
│ id  ┆ a   ┆ b   │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═════╪═════╪═════╡
│ x   ┆ NaN ┆ 5.0 │
│ y   ┆ 3.5 ┆ 7.5 │
└─────┴─────┴─────┘

我希望将加权平均值中的 NaN 替换为等加权平均值中的 2，因为当忽略“a”列中空值的权重时，权重之和为 0，因此等加权平均值应该被退回。

使用 pandas 我可以通过以下方式完成此任务

import pandas

def _wavg_py(gr, vr, wt):
    x = gr[[vr, wt]].dropna()
    den = x[wt].sum()
    if(den == 0):
        return(gr[vr].mean())
    else:
        return((x[vr] * x[wt]).sum() / den)

res = df.groupby(by = gr).apply(_wavg_py, col, wt)

Answer 1

我不确定这是否是可以接受的方法

.fill_nan()

与

.mean()

df.group_by(gr, maintain_order=True).agg(
    wavg.fill_nan(pl.col(vr).mean())
)

shape: (2, 3)
┌─────┬─────┬─────┐
│ id  ┆ a   ┆ b   │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═════╪═════╪═════╡
│ x   ┆ 2.0 ┆ 5.0 │
│ y   ┆ 3.5 ┆ 7.5 │
└─────┴─────┴─────┘

或者，您似乎在问如何在表达式内执行条件逻辑

即

when/then/otherwise

df.group_by(gr, maintain_order=True).agg(
    pl.when(pl.col('wt').filter(pl.col(vr).is_not_null()).sum() == 0)
      .then(pl.col(vr).mean())
      .otherwise(wavg)
)

shape: (2, 3)
┌─────┬─────┬─────┐
│ id  ┆ a   ┆ b   │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═════╪═════╪═════╡
│ x   ┆ 2.0 ┆ 5.0 │
│ y   ┆ 3.5 ┆ 7.5 │
└─────┴─────┴─────┘