使用布尔掩码和不同大小的数组进行极坐标算术运算

Question

我非常努力地尝试使用ChatGPT，但无法完成这项工作，也无法解决它。我应该如何使用布尔掩码分配或增加一组值？这就是我对熊猫所做的事情。

validation_predictions = model.predict(validation_dataset)
df.loc[validation_mask, 'prediction'] += validation_predictions / len(seeds)

validation_mask

有 N 个

True

，并且

validation_predictions

是一个大小为

N

的 numpy 数组，所以这样的赋值工作正常。然而我无法用极地达到同样的效果。

我尝试了when/then/otherwise链，但它会抛出错误，因为

validation_predictions

大小与整个数据帧的大小不匹配。

df = df.with_columns(
    pl.when(validation_mask)
    .then(pl.col('prediction') + (validation_predictions / len(seeds)))
    .otherwise(pl.col('prediction'))
    .alias('prediction')
)
# ShapeError: cannot evaluate two Series of different lengths (6 and 5)

为了重现性目的：

import polars as pl
import numpy as np

seeds = [42, 1337, 0]

df = pl.DataFrame({
    "some_column": [10, 20, 30, 40, 50, 60],
    "prediction": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
})

validation_mask = df["some_column"] > 15
validation_predictions = np.array([0.5, 1.5, 2.5, 3.5, 4.5])

预期结果：

shape: (6, 2)
┌─────────────┬────────────┐
│ some_column ┆ prediction │
│ ---         ┆ ---        │
│ i64         ┆ f64        │
╞═════════════╪════════════╡
│ 10          ┆ 1.0        │
│ 20          ┆ 2.166667   │
│ 30          ┆ 3.5        │
│ 40          ┆ 4.833333   │
│ 50          ┆ 6.166667   │
│ 60          ┆ 7.5        │
└─────────────┴────────────┘

Answer 1

使用

df.with_row_index

获取原始位置，根据该位置进行过滤，然后执行 Join。

如果您使用内部联接（默认），它将删除过滤掉的行，但您可以使用左联接，然后使用

fill_null

进行合并

import polars as pl
import numpy as np

df = pl.DataFrame({
    "some_column": [10, 20, 30, 40, 50, 60],
    "prediction": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
})
validation_mask = df["some_column"] > 15
validation_predictions = np.array([0.5, 1.5, 2.5, 3.5, 4.5])

df = df.with_row_index('idx')

positions = df.select(pl.col("idx").filter(validation_mask))

val_predictions = positions.with_columns(val_prediction = validation_predictions)

joined = df.join(val_predictions, on="idx", how="left")
result = joined.with_columns(pl.col("prediction") + pl.col("val_prediction").fill_null(0)).drop('val_prediction')

print(result)