如何（更快）找到使涉及多个 LazyFrame 列的函数的结果最小化的 x 值？

Question

我有一个 LazyFrame，其中包含几个时期内的多列每小时数据。对于每个周期，我想找到涉及多列数学运算的函数的 x 值，以最小化结果。

\n

我使用 scipy.optimize.minimize 来完成此操作，并且实际上获得了所需的结果。问题是这个过程运行得非常慢，所以我只是在寻找任何能完成相同但更快的事情。

\n

    def minimization_target(x, period_start):\n        return hourly_data.filter(pl.col('period_start') == period_start).select((((pl.col('price').median() * pl.col('quantity').median() - (pl.col('estimated_quantity') * (pl.col('estimated_price') + x)).sum()) / (pl.col('key_product') * (pl.col('estimated_price') + x)).sum())).abs() - 1).abs()).collect().item()\n\n    results = hourly_data.group_by('period_start', maintain_order=True).map_groups(lambda group: pl.DataFrame({'x_values': scipy.optimize.minimize(minimization_target, group.get_column('initial_guess').median(), args=group.get_column('period_start').median()).x}), schema=None)\n

\n

最小的例子：

\n

import scipy\nimport polars as pl\nfrom datetime import datetime\n\nhourly_data = pl.DataFrame({'period': [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3], 'price': [4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7], 'quantity': [7, 7, 7, 7, 7, 7, 6, 6, 6, 6, 6, 6, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4], 'estimated_price': [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8], 'estimated_quantity': [6, 6, 6, 6, 6, 6, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3], 'key_product': [0.9, 0.8, 0.7, 0.8, 0.9, 0.8, 0.7, 0.8, 0.9, 0.8, 0.7, 0.8, 0.9, 0.8, 0.7, 0.8, 0.9, 0.8, 0.7, 0.8, 0.9, 0.8, 0.7, 0.8], 'initial_guess': [10, 10, 10, 10, 10, 10, 20, 20, 20, 20, 20, 20, 30, 30, 30, 30, 30, 30, 40, 40, 40, 40, 40, 40]}).lazy()\nhourly_data = hourly_data.with_columns(pl.datetime_range(datetime(2024, 1, 1), datetime(2024, 1, 1, 23), '1h').alias('hour'))\nhourly_data = hourly_data.with_columns(pl.col('hour').min().over('period').alias('period_start'))\n\n\ndef minimization_target(x, period_start):\n    return hourly_data.filter(pl.col('period_start') == period_start).select((((pl.col('price').median() * pl.col('quantity').median() - (pl.col('estimated_quantity') * (pl.col('estimated_price') + x)).sum()) / (pl.col('key_product') * (pl.col('estimated_price') + x)).sum()).abs() - 1).abs()).collect().item()\n\n\nresults = hourly_data.group_by('period_start', maintain_order=True).map_groups(lambda group: pl.DataFrame({'x_values': scipy.optimize.minimize(minimization_target, group.get_column('initial_guess').median(), args=group.get_column('period_start').median()).x}), schema=None)\n

\n

输入：

\n

shape: (24, 9)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 period \xe2\x94\x86 price \xe2\x94\x86 quantity \xe2\x94\x86 estimated_p \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 key_product \xe2\x94\x86 initial_gu \xe2\x94\x86 hour       \xe2\x94\x86 period_sta \xe2\x94\x82\n\xe2\x94\x82 ---    \xe2\x94\x86 ---   \xe2\x94\x86 ---      \xe2\x94\x86 rice        \xe2\x94\x86   \xe2\x94\x86 ---         \xe2\x94\x86 ess        \xe2\x94\x86 ---        \xe2\x94\x86 rt         \xe2\x94\x82\n\xe2\x94\x82 i64    \xe2\x94\x86 i64   \xe2\x94\x86 i64      \xe2\x94\x86 ---         \xe2\x94\x86   \xe2\x94\x86 f64         \xe2\x94\x86 ---        \xe2\x94\x86 datetime[\xce\xbc \xe2\x94\x86 ---        \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86 i64         \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86 i64        \xe2\x94\x86 s]         \xe2\x94\x86 datetime[\xce\xbc \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86            \xe2\x94\x86 s]         \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 0      \xe2\x94\x86 4     \xe2\x94\x86 7        \xe2\x94\x86 5           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.9         \xe2\x94\x86 10         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 00:00:00   \xe2\x94\x86 00:00:00   \xe2\x94\x82\n\xe2\x94\x82 0      \xe2\x94\x86 4     \xe2\x94\x86 7        \xe2\x94\x86 5           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.8         \xe2\x94\x86 10         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 01:00:00   \xe2\x94\x86 00:00:00   \xe2\x94\x82\n\xe2\x94\x82 0      \xe2\x94\x86 4     \xe2\x94\x86 7        \xe2\x94\x86 5           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.7         \xe2\x94\x86 10         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 02:00:00   \xe2\x94\x86 00:00:00   \xe2\x94\x82\n\xe2\x94\x82 0      \xe2\x94\x86 4     \xe2\x94\x86 7        \xe2\x94\x86 5           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.8         \xe2\x94\x86 10         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 03:00:00   \xe2\x94\x86 00:00:00   \xe2\x94\x82\n\xe2\x94\x82 0      \xe2\x94\x86 4     \xe2\x94\x86 7        \xe2\x94\x86 5           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.9         \xe2\x94\x86 10         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 04:00:00   \xe2\x94\x86 00:00:00   \xe2\x94\x82\n\xe2\x94\x82 \xe2\x80\xa6      \xe2\x94\x86 \xe2\x80\xa6     \xe2\x94\x86 \xe2\x80\xa6        \xe2\x94\x86 \xe2\x80\xa6           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 \xe2\x80\xa6           \xe2\x94\x86 \xe2\x80\xa6          \xe2\x94\x86 \xe2\x80\xa6          \xe2\x94\x86 \xe2\x80\xa6          \xe2\x94\x82\n\xe2\x94\x82 3      \xe2\x94\x86 7     \xe2\x94\x86 4        \xe2\x94\x86 8           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.8         \xe2\x94\x86 40         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 19:00:00   \xe2\x94\x86 18:00:00   \xe2\x94\x82\n\xe2\x94\x82 3      \xe2\x94\x86 7     \xe2\x94\x86 4        \xe2\x94\x86 8           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.9         \xe2\x94\x86 40         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 20:00:00   \xe2\x94\x86 18:00:00   \xe2\x94\x82\n\xe2\x94\x82 3      \xe2\x94\x86 7     \xe2\x94\x86 4        \xe2\x94\x86 8           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.8         \xe2\x94\x86 40         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 21:00:00   \xe2\x94\x86 18:00:00   \xe2\x94\x82\n\xe2\x94\x82 3      \xe2\x94\x86 7     \xe2\x94\x86 4        \xe2\x94\x86 8           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.7         \xe2\x94\x86 40         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 22:00:00   \xe2\x94\x86 18:00:00   \xe2\x94\x82\n\xe2\x94\x82 3      \xe2\x94\x86 7     \xe2\x94\x86 4        \xe2\x94\x86 8           \xe2\x94\x86 \xe2\x80\xa6 \xe2\x94\x86 0.8         \xe2\x94\x86 40         \xe2\x94\x86 2024-01-01 \xe2\x94\x86 2024-01-01 \xe2\x94\x82\n\xe2\x94\x82        \xe2\x94\x86       \xe2\x94\x86          \xe2\x94\x86             \xe2\x94\x86   \xe2\x94\x86             \xe2\x94\x86            \xe2\x94\x86 23:00:00   \xe2\x94\x86 18:00:00   \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

期望的输出：

\n

shape: (4, 1)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 x_values   \xe2\x94\x82\n\xe2\x94\x82 ---        \xe2\x94\x82\n\xe2\x94\x82 f64        \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 -16.006287 \xe2\x94\x82\n\xe2\x94\x82 10.331055  \xe2\x94\x82\n\xe2\x94\x82 25.420471  \xe2\x94\x82\n\xe2\x94\x82 37.352234  \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

Answer 1

避免使用整个数据框。

首先，您可以仅使用优化目标中的当前组而不是首先过滤整个数据帧，从而将运行时间缩短约 2 倍。

这可以通过如下方式实现。

import functools

def minimization_target(x, group: pl.DataFrame):
    numerator_expr = pl.col('price').median() * pl.col('quantity').median() - (pl.col('estimated_quantity') * (pl.col('estimated_price') + x)).sum()
    denominator_expr = (pl.col('key_product') * (pl.col('estimated_price') + x)).sum()
    return (
        group.select(
            (
                (numerator_expr / denominator_expr).abs() - 1
            ).abs()
        )
        .item()
    )

(
    hourly_data
    .group_by('period_start', maintain_order=True)
    .map_groups(
        lambda group: pl.DataFrame({"x": scipy.optimize.minimize(functools.partial(minimization_target, group=group), group.get_column('initial_guess').median()).x}),
        schema=None,
    )
    .collect()
)

简化最小化目标。

我们希望最小化目标尽可能简单，以便scipy.optimize.minimize重复评估函数。当前的实现适用于“重”DataFrame 对象，并计算所有评估通用的许多术语。

在预先计算这些项的同时使用更简单的最小化目标可进一步提高约 4 倍的速度。

import numpy as np

def minimization_target(x: float, s1: float, f1: float, s2: float, f2: float):
    return np.abs(np.abs((s1 - f1 * x) / (s2 + f2 * x)) - 1)

(
    hourly_data
    .group_by('period_start', maintain_order=True)
    .map_groups(
        lambda group: pl.DataFrame({
            "x": scipy.optimize.minimize(
                fun=partial(
                    minimization_target,
                    # precomputing terms common to all evaluations of the minimization target
                    s1=group.select(pl.col('price').median() * pl.col('quantity').median() - (pl.col('estimated_quantity') * pl.col('estimated_price')).sum()).item(),
                    f1=group.select(pl.col("estimated_quantity").sum()).item(),
                    s2=group.select((pl.col('key_product') * pl.col('estimated_price')).sum()).item(),
                    f2=group.select(pl.col("key_product").sum()).item(),
                ),
                x0=group.get_column('initial_guess').median()
            ).x
        }),
        schema=None,
    )
    .collect()
)