使用 Polars 生成运行记录，其中一个月的期末余额将成为下个月的期初余额

我想在给定初始期初余额的情况下生成一段时间内的运行计数模拟。例如，一个虚构的企业从 1000 美元开始，在任何一个月内都会经历随机的盈利和亏损。 12 个月后，这家企业可能会赚多少钱？ 我的问题是，我可以生成所需的列，但无法弄清楚如何在不破坏当月数据的情况下将模拟计数延续到下个月。

我可以创建一个 Polars 数据框，并用所需的损益列填充该数据框，并生成期末余额。 使用堆叠的“with_columns()”，我为期初余额和期末余额创建了所需的列，然后生成一些随机利润和损失，然后更新期末余额。但我无法让这件事继续下去。

import polars as pl
import datetime as dt
from dateutil.relativedelta import relativedelta
start_date = dt.date.today() + relativedelta(months=1, day=1)
df = pl.DataFrame(
    pl.date_range(start_date, start_date + relativedelta(months=5), '1mo').alias('date'),
)
beginning_balance = 1000.0
df = df.with_columns(
        pl.lit(beginning_balance).alias('beginning_balance'),
        pl.lit(beginning_balance).alias('closing_balance'),
).with_columns(
    [
        pl.when(pl.col('date') == start_date)
        .then(pl.col('beginning_balance'))
        .otherwise(pl.col('closing_balance').shift(1))
        .alias('beginning_balance'), 
        pl.Series([normalvariate(100, 80) for _ in range(len(df))]).round(2).alias('profit'),
        pl.Series([normalvariate(100, 75) for _ in range(len(df))]).round(2).alias('loss'),
    ]
).with_columns(
    (pl.col('beginning_balance') + pl.col('profit') - pl.col('loss')).alias('closing_balance'),
)
df

导致 ：

 shape: (6, 5)
date    beginning_balance   closing_balance     profit  loss
date    f64     f64     f64     f64
2023-04-01  1000.0  909.61  72.97   163.36
2023-05-01  1000.0  1072.29     212.84  140.55
2023-06-01  1000.0  974.54  82.75   108.21
2023-07-01  1000.0  824.08  -60.69  115.23
2023-08-01  1000.0  940.71  155.6   214.89
2023-09-01  1000.0  994.5   94.23   99.73

请注意，第一个月的期末余额 (909.61) 并未转入下个月，下个月仍为 1000.00，因此不会发生运行计数。

我绝对可以使用循环（for 或 while）来遍历每一行，但这会阻碍 Polars 的优化和并行化。以下方法可以工作，但速度要慢得多（特别是当我运行许多具有更详细计算的模拟时）。

current_tally = beginning_balance
for t in range(len(df)):
    beginning_balance = current_tally
    current_tally = beginning_balance + df[t, 'profit'] - df[t, 'loss']
    df[t, 'beginning_balance'] = beginning_balance
    df[t, 'closing_balance'] = current_tally

df
 shape: (6, 5)
date    beginning_balance   closing_balance     profit  loss
date    f64     f64     f64     f64
2023-04-01  1000.0  1063.13     115.49  52.36
2023-05-01  1063.13     1224.37     184.62  23.38
2023-06-01  1224.37     1258.02     114.68  81.03
2023-07-01  1258.02     1323.73     239.12  173.41
2023-08-01  1323.73     1341.33     102.18  84.58
2023-09-01  1341.33     1329.88     19.9    31.35

我该怎么做within Polars?

我认为替代解决方案更有效（因为.cumulative_eval多次进行大量添加）：

开始数据框date和随机的profit and loss:

import polars as pl
import datetime as dt
from dateutil.relativedelta import relativedelta

num_months = 6
start_date = dt.date.today() + relativedelta(months=1, day=1)
df = pl.DataFrame().with_columns([
    pl.date_range(start_date, start_date + relativedelta(months=num_months - 1), '1mo').alias('date'),
    pl.Series([normalvariate(100, 80) for _ in range(num_months)]).round(2).alias('profit'),
    pl.Series([normalvariate(100, 75) for _ in range(num_months)]).round(2).alias('loss'),
])

基于列构建新列profit and loss:

beginning_balance = 1000.0
df = df.with_columns(
    (beginning_balance + (pl.col('profit') - pl.col('loss')).cumsum())
    .alias('closing_balance')
).with_columns(
    pl.col('closing_balance').shift_and_fill(1, beginning_balance)
    .alias('beginning_balance')
)

The tally通过使用计算.cumsum()关于之间的区别profit and loss.