Zeilenweises Skalarprodukt in Polaren ⇐ Python

[Anonymous]

Ich habe einen Datenrahmen mit zwei Spaltenwerten und Gewichtungen vom Typ list[i64] und möchte ein zeilenweises Skalarprodukt der beiden durchführen.

Code: Select all

df = pl.DataFrame({
'values': [[0], [0, 2], [0, 2, 4], [2, 4, 0], [4, 0, 8]],
'weights': [[3], [2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]
})

Es gibt eine Möglichkeit, die funktioniert hat: Zuerst Werte und Gewichte in eine Struktur einfügen und dann .map_elements für jede Zeile ausführen:

Code: Select all

df.with_columns(
pl.struct(['values', 'weights'])
.map_elements(
lambda x: np.dot(x['values'], x['weights']), return_dtype=pl.Float64
).alias('dot')
)

Aber wie aus der Dokumentation hervorgeht, ist map_elements im Allgemeinen viel langsamer als native Polarausdrücke, daher habe ich versucht, es in nativen Ausdrücken zu implementieren.
Ich habe Folgendes versucht:

Code: Select all

df.with_columns(
pl.concat_list('values', 'weights').alias('combined'),
pl.concat_list('values', 'weights').list.eval(pl.element().slice(0, pl.len() // 2)).alias('values1'),
pl.concat_list('values', 'weights').list.eval(pl.element().slice(pl.len() // 2, pl.len() // 2)).alias('values2'),
pl.concat_list('values', 'weights').list.eval(
pl.element().slice(0, pl.len() // 2).dot(pl.element().slice(pl.len() // 2, pl.len() // 2))
).list.first().alias('dot'),
pl.concat_list('values', 'weights').list.eval(
pl.element().slice(0, pl.len() // 2) + pl.element().slice(pl.len() // 2, pl.len() // 2)
).alias('sum'),
)

Ich habe erwartet, dass die Punktspalte [0, 6, 16, 10, 28] ist, aber es stellt sich heraus, dass es das Folgende ist.

Code: Select all

shape: (5, 7)
┌───────────┬───────────┬─────────────┬───────────┬───────────┬─────┬────────────┐
│ values    ┆ weights   ┆ combined    ┆ values1   ┆ values2   ┆ dot ┆ sum        │
│ ---       ┆ ---       ┆ ---         ┆ ---       ┆ ---       ┆ --- ┆ ---        │
│ list[i64] ┆ list[i64] ┆ list[i64]   ┆ list[i64] ┆ list[i64] ┆ i64 ┆ list[i64]  │
╞═══════════╪═══════════╪═════════════╪═══════════╪═══════════╪═════╪════════════╡
│ [0]       ┆ [3]       ┆ [0, 3]      ┆ [0]       ┆ [3]       ┆ 0   ┆ [0]        │
│ [0, 2]    ┆ [2, 3]    ┆ [0, 2, … 3] ┆ [0, 2]    ┆ [2, 3]    ┆ 4   ┆ [0, 4]     │
│ [0, 2, 4] ┆ [1, 2, 3] ┆ [0, 2, … 3] ┆ [0, 2, 4] ┆ [1, 2, 3] ┆ 20  ┆ [0, 4, 8]  │
│ [2, 4, 0] ┆ [1, 2, 3] ┆ [2, 4, … 3] ┆ [2, 4, 0] ┆ [1, 2, 3] ┆ 20  ┆ [4, 8, 0]  │
│ [4, 0, 8] ┆ [1, 2, 3] ┆ [4, 0, … 3] ┆ [4, 0, 8] ┆ [1, 2, 3] ┆ 80  ┆ [8, 0, 16] │
└───────────┴───────────┴─────────────┴───────────┴───────────┴─────┴────────────┘

Beachten Sie, dass nicht einmal die Summe meinen Erwartungen entspricht. Das erste Slice scheint sich selbst hinzuzufügen und nicht das zweite Slice
Mache ich etwas falsch? Was ist der beste Weg, um ein zeilenweises Punktprodukt in Polars durchzuführen?