如何在 Polars DataFrame 中创建某些列的加权和？

Question

我正在寻找一种惯用的方法来计算 Polars DataFrame 中列子集的加权和，并将其作为新列添加到 DataFrame 中。因此，假设我想将下面 DataFrame 中的列 p1-p3 乘以以下权重，然后将它们相加以创建一个新列。

weights = [7.4, 3.2, -0.13]

df = pl.DataFrame(
    {
        "id": [1, 2, 3, 4],
        "p1": [44.3, 2.3, 2.4, 6.2],
        "p2": [7.3, 8.4, 10.3, 8.443],
        "p3": [70.3, 80.4, 100.3, 80.443],
        "p4": [16.4, 18.2, 11.5, 18.34],
    }
)
df

id   p1     p2      p3      p4
i64  f64    f64     f64     f64
1    44.3   7.3     70.3    16.4
2    2.3    8.4     80.4    18.2
3    2.4    10.3    100.3   11.5
4    6.2    8.443   80.443  18.34

我想出了以下计算正确答案的解决方案，但我觉得可能有一种更简单、更惯用的方法可以让我 select 感兴趣的列，而无需重新- 在 with_columns 函数中指定 df。有什么建议吗？

df.with_columns(
    [
        df.select(
            [
                pl.col(col) * pl.lit(weights[i])
                for i, col in enumerate(["p1", "p2", "p3"])
            ]
        )
        .fold(lambda c1, c2: c1 + c2)
        .alias("index"),
    ]
)

id  p1    p2     p3      p4     index
i64 f64   f64    f64     f64    f64
1   44.3  7.3    70.3    16.4   342.041
2   2.3   8.4    80.4    18.2   33.448
3   2.4   10.3   100.3   11.5   37.681
4   6.2   8.443  80.443  18.34  62.44

Answer 1

你快到了。可以直接使用pl.fold表达式。


df.with_columns([
    pl.fold(acc=0, f=lambda c1, c2: c1 + c2, exprs=[
        pl.col(col) * pl.lit(weights[i])
        for i, col in enumerate(["p1", "p2", "p3"])
    ]).alias("index")
])

shape: (4, 6)
┌─────┬──────┬───────┬────────┬───────┬─────────┐
│ id  ┆ p1   ┆ p2    ┆ p3     ┆ p4    ┆ index   │
│ --- ┆ ---  ┆ ---   ┆ ---    ┆ ---   ┆ ---     │
│ i64 ┆ f64  ┆ f64   ┆ f64    ┆ f64   ┆ f64     │
╞═════╪══════╪═══════╪════════╪═══════╪═════════╡
│ 1   ┆ 44.3 ┆ 7.3   ┆ 70.3   ┆ 16.4  ┆ 342.041 │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2.3  ┆ 8.4   ┆ 80.4   ┆ 18.2  ┆ 33.448  │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ 2.4  ┆ 10.3  ┆ 100.3  ┆ 11.5  ┆ 37.681  │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 4   ┆ 6.2  ┆ 8.443 ┆ 80.443 ┆ 18.34 ┆ 62.44   │
└─────┴──────┴───────┴────────┴───────┴─────────┘

Answer 2

一个可能也有帮助的技巧：sum、max、min、all 和 any 函数中内置了常用的折叠- 当使用 表达式列表 .

调用时

例如：

pl.sum([exp1, exp2, etc...])

是常用加法折叠的语法糖：

pl.fold(pl.lit(0), f=lambda c1, c2: c1 + c2, exprs =[expr1, expr2, etc...])

因此，我们可以做到以下几点：

col_names = ["p1", "p2", "p3"]
weights = [7.4, 3.2, -0.13]
df.with_column(
    pl.sum(
        [pl.col(col_nm) * wgt
         for col_nm, wgt in zip(col_names, weights)]
    ).alias("index")
)

shape: (4, 6)
┌─────┬──────┬───────┬────────┬───────┬─────────┐
│ id  ┆ p1   ┆ p2    ┆ p3     ┆ p4    ┆ index   │
│ --- ┆ ---  ┆ ---   ┆ ---    ┆ ---   ┆ ---     │
│ i64 ┆ f64  ┆ f64   ┆ f64    ┆ f64   ┆ f64     │
╞═════╪══════╪═══════╪════════╪═══════╪═════════╡
│ 1   ┆ 44.3 ┆ 7.3   ┆ 70.3   ┆ 16.4  ┆ 342.041 │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2.3  ┆ 8.4   ┆ 80.4   ┆ 18.2  ┆ 33.448  │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ 2.4  ┆ 10.3  ┆ 100.3  ┆ 11.5  ┆ 37.681  │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 4   ┆ 6.2  ┆ 8.443 ┆ 80.443 ┆ 18.34 ┆ 62.44   │
└─────┴──────┴───────┴────────┴───────┴─────────┘

我使用 zip 而不是 enumerate .. 但这是一种风格选择。我允许 Polars 将权重标量广播为文字，而不是显式使用 pl.lit.

另一个可能有助于提高可读性的技巧：我们可以在 with_column/with_columns/select 上下文之外生成表达式列表。

例如：

col_names = ["p1", "p2", "p3", "p4"]
weights = [7.4, 3.2, -0.13, 0.0]
wghtd_cols = [
    pl.col(col_nm) * wgt
    for col_nm, wgt in zip(col_names, weights)
    if wgt != 0.0
]

df.with_column(pl.sum(wghtd_cols).alias("index"))

shape: (4, 6)
┌─────┬──────┬───────┬────────┬───────┬─────────┐
│ id  ┆ p1   ┆ p2    ┆ p3     ┆ p4    ┆ index   │
│ --- ┆ ---  ┆ ---   ┆ ---    ┆ ---   ┆ ---     │
│ i64 ┆ f64  ┆ f64   ┆ f64    ┆ f64   ┆ f64     │
╞═════╪══════╪═══════╪════════╪═══════╪═════════╡
│ 1   ┆ 44.3 ┆ 7.3   ┆ 70.3   ┆ 16.4  ┆ 342.041 │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2.3  ┆ 8.4   ┆ 80.4   ┆ 18.2  ┆ 33.448  │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ 2.4  ┆ 10.3  ┆ 100.3  ┆ 11.5  ┆ 37.681  │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 4   ┆ 6.2  ┆ 8.443 ┆ 80.443 ┆ 18.34 ┆ 62.44   │
└─────┴──────┴───────┴────────┴───────┴─────────┘

当您的代码的一部分正在生成权重 and/or 选择列，而您的代码的另一部分正在 DataFrame 中创建结果加权和列时，这特别有用。

如何在 Polars DataFrame 中创建某些列的加权和？

How to create a weighted sum of some columns in a Polars DataFrame?

python-polars