Я ищу идиоматический способ вычислить взвешенную сумму подмножества столбцов в фрейме данных Polars и добавить его в фрейм данных в качестве нового столбца. Итак, допустим, я хочу умножить столбцы p1-p3 в DataFrame ниже на следующие веса, а затем суммировать их, чтобы создать новый столбец.
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
Я придумал следующее решение, которое вычисляет правильный ответ, но я чувствую, что, вероятно, существует более простой и идиоматический метод, который позволил бы мне выбрать интересующие столбцы без необходимости повторного указания df
в функции with_columns
. Какие-либо предложения?
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
Ты был почти там. Вы можете использовать выражение 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 │
└─────┴──────┴───────┴────────┴───────┴─────────┘
Один трюк, который также может помочь: часто используемые складки встроены в функции 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 │
└─────┴──────┴───────┴────────┴───────┴─────────┘
Это особенно полезно, когда одна часть вашего кода генерирует веса и/или выбирает столбцы, а другая часть вашего кода создает результирующий столбец взвешенной суммы в DataFrame.
Хороший ответ. Спасибо, что подняли это на новый уровень и предоставили альтернативные составы.
Да, это то место, куда я думал, что я должен быть в состоянии добраться. Спасибо,