如何用Python和Polars实现带利息与月结余额结转的储蓄模拟

核心解决方案 #

利用Polars的cumulative_eval实现递推式的期末余额计算，完全替代Python循环，充分发挥Polars的向量化性能优势。以下是完整实现：

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

# 配置参数
starting_balance = 500.00
milestone = 6
income_average = 700.00
income_std = 300.00
interest_average = 0.10
interest_std = 0.10

# 生成日期序列
start_date = dt.date.today().replace(day=1)
date_range = pl.date_range(
    start=start_date,
    end=start_date + relativedelta(months=milestone),
    interval="1mo",
    eager=True
)

# 预生成所有月份的利息和收入数据（首行对应初始状态，值为0）
interest_vals = [0.0] + [normalvariate(interest_average, interest_std) for _ in range(milestone)]
income_vals = [0.0] + [round(normalvariate(income_average, income_std), 2) for _ in range(milestone)]

# 构建DataFrame并完成所有计算
df = pl.DataFrame({
    "date": date_range,
    "interest": interest_vals,
    "income": income_vals
}).with_columns(
    # 递推计算期末余额：当月余额 = 上月余额*(1+当月利息) + 当月收入
    closing_balance=pl.cumulative_eval(
        lambda prev, curr: prev * (1 + curr["interest"]) + curr["income"],
        initial=pl.lit(starting_balance)
    ).round(2),
    # 期初余额：首行为0，后续为上月期末余额
    opening_balance=pl.col("closing_balance").shift(1).fill_null(0.0)
).select(
    # 调整列顺序与原循环版本一致
    "date", "opening_balance", "interest", "income", "closing_balance"
)

print(df)

关键细节说明 #

1. 预生成随机数据：一次性生成所有月份的利息和收入，避免循环中逐行生成的低效操作
2. cumulative_eval递推逻辑：
   • prev代表上一行的期末余额
   • curr代表当前行的利息和收入数据
   • 完全贴合业务公式：期末余额 = 期初余额*(1+利息) + 收入
3. 期初余额生成：通过shift(1)将期末余额向下偏移一行，首行用fill_null(0.0)填充，和原循环逻辑完全对齐
4. 精度控制：用round(2)保证金额保留两位小数，符合财务数据规范

性能优势 #

• 完全基于Polars向量化运算，摆脱Python循环的性能瓶颈，数据量越大（如模拟几十年退休储蓄），性能提升越显著
• 代码更简洁，逻辑更清晰，符合Polars声明式编程风格
• 支持Lazy API进一步优化（将eager=True移除，最后加.collect()），适合大规模数据模拟

输出示例 #

运行后输出格式与原循环版本完全一致：

┌────────────┬─────────────────┬──────────┬─────────┬─────────────────┐
│ date       ┆ opening_balance ┆ interest ┆ income  ┆ closing_balance │
│ ---        ┆ ---             ┆ ---      ┆ ---     ┆ ---             │
│ date       ┆ f64             ┆ f64      ┆ f64     ┆ f64             │
╞════════════╪═════════════════╪══════════╪═════════╪═════════════════╡
│ 2024-05-01 ┆ 0.0             ┆ 0.0      ┆ 0.0     ┆ 500.0           │
│ 2024-06-01 ┆ 500.0           ┆ 0.1875   ┆ 723.45  ┆ 1317.20         │
│ 2024-07-01 ┆ 1317.20         ┆ 0.0523   ┆ 489.12  ┆ 1874.18         │
│ ...        ┆ ...             ┆ ...      ┆ ...     ┆ ...             │
└────────────┴─────────────────┴──────────┴─────────┴─────────────────┘

内容的提问来源于stack exchange，提问作者Buckley