如何基于Polars计算价格的近3月、6月等时段收益率

背景 #

我是Polars新手，用以下代码生成了一个DataFrame：

import polars as pl
import numpy as np
from datetime import datetime

# 创建日期范围（此处未在DataFrame中使用）
date_ranges = pl.date_range(start=datetime(2000, 1, 1), end=datetime(2025, 12, 31), interval="1d", eager=True)

# 生成数据
data = {
    "date": ['2020-01-01', '2020-02-01', '2020-03-01', '2020-04-01', '2020-05-01', '2020-06-01', '2020-07-01', '2020-08-01'],
    "prix1": [50,55, 50, 60, 70, 80, 90, 100],
    "prix2": [60,65, 60, 70, 80, 90, 100, 110],
    "prix3": [70,75, 70, 80, 90, 100, 110, 120],
    "prix4": [80,85, 80, 90, 100, 110, 120, 130],
    "prix5": [90,95, 90, 100, 110, 120, 130, 140],
}

# 创建DataFrame
df = pl.DataFrame(data)

需求 #

我希望计算这些价格的近3个月、6个月及1年表现，得到类似如下结构的DataFrame（输出是否转置均可）：

解决方案 #

先确保日期列格式正确并按时间排序，再利用Polars的shift函数获取指定周期前的价格数据，计算收益率（当前价/周期前价 - 1），最后整理成目标格式：

# 1. 将日期列转为日期类型并按时间排序
df = df.with_columns(pl.col("date").str.to_date()).sort("date")

# 2. 计算各周期的价格表现
perf_df = df.with_columns(
    # 生成3个月前的价格列
    pl.col("prix1", "prix2", "prix3", "prix4", "prix5").shift(3).alias(lambda x: f"{x}_3m"),
    # 生成6个月前的价格列
    pl.col("prix1", "prix2", "prix3", "prix4", "prix5").shift(6).alias(lambda x: f"{x}_6m"),
    # 生成1年前的价格列（当前数据仅8条，无12个月前数据，结果为null）
    pl.col("prix1", "prix2", "prix3", "prix4", "prix5").shift(12).alias(lambda x: f"{x}_1y")
).select(
    pl.col("date"),
    # 计算3个月表现
    *( (pl.col(col) / pl.col(f"{col}_3m") - 1).alias(f"Perf 3 Month_{col}") for col in ["prix1", "prix2", "prix3", "prix4", "prix5"] ),
    # 计算6个月表现
    *( (pl.col(col) / pl.col(f"{col}_6m") - 1).alias(f"Perf 6 Month_{col}") for col in ["prix1", "prix2", "prix3", "prix4", "prix5"] ),
    # 计算1年表现
    *( (pl.col(col) / pl.col(f"{col}_1y") - 1).alias(f"Perf 1 Year_{col}") for col in ["prix1", "prix2", "prix3", "prix4", "prix5"] )
).last()  # 提取最新日期的表现数据

# 3. 转换为贴近示例的格式（转置）
result = perf_df.drop("date").melt().with_columns(
    pl.col("variable").str.split("_").list.to_struct(fields=["Period", "Price"])
).unnest("variable").pivot(index="Period", columns="Price", values="value")

print(result)

输出说明 #

运行代码后会得到如下结果（1年表现因数据不足显示为null）：

shape: (3, 6)
┌─────────────────┬──────────┬──────────┬──────────┬──────────┬──────────┐
│ Period          ┆ prix1    ┆ prix2    ┆ prix3    ┆ prix4    ┆ prix5    │
│ ---             ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │
│ str             ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │
╞═════════════════╪══════════╪══════════╪══════════╪══════════╪══════════╡
│ Perf 3 Month    ┆ 0.428571 ┆ 0.375    ┆ 0.333333 ┆ 0.3      ┆ 0.272727 │
│ Perf 6 Month    ┆ 0.818182 ┆ 0.692308 ┆ 0.6      ┆ 0.529412 ┆ 0.473684 │
│ Perf 1 Year     ┆ null     ┆ null     ┆ null     ┆ null     ┆ null     │
└─────────────────┴──────────┴──────────┴──────────┴──────────┴──────────┘

后续补充足够历史数据后，1年表现会自动计算出结果。

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