如何在Python的DataFrame中计算指定周期的股票复利收益？

我有一个包含日期、股票价格、股息列的DataFrame，想要新增一列，按照指定公式计算复利收益。当n=2时，需使用当前行及前两行数据计算收益，示例DataFrame如下：

price  dividend   
2020-07-31  83.08    0.7125  
2020-08-31  73.35    0.7225 
2020-09-30  74.55    0.7325  
2020-10-31  81.57    0.8400  
2020-11-30  81.85    0.8500  
2020-12-31  79.95    0.8600

以2020-12-31行为例，计算逻辑为：
CompoundedReturn = (79.95 / 81.57) * ((1 + 0.84/81.57) * (1 + 0.85/81.85) * (1 + 0.86/79.95)) = 1.0113
计算后的目标DataFrame如下：

price  dividend    return
2020-07-31  83.08    0.7125        NA     
2020-08-31  73.35    0.7225        NA     
2020-09-30  74.55    0.7325    0.9229     
2020-10-31  81.57    0.8400    1.1457     
2020-11-30  81.85    0.8500    1.1318     
2020-12-31  79.95    0.8600    1.0113

请问Python或numpy中是否有可实现此需求的内置函数？

Python或NumPy中没有直接对应这个特定复利收益计算的内置函数，但可以通过Pandas的窗口函数和向量化操作高效实现，无需逐行循环：

• 拆解公式逻辑
  公式可拆分为两部分：

  1. 价格比例项：当前行价格 / n周期前的价格
  2. 股息复利项：n+1个周期（从n周期前到当前行）的(1 + 股息/价格)的乘积

• 代码实现
  假设你的DataFrame名为df，n=2，代码如下：

  import pandas as pd
  
  # 计算每个周期的(1 + dividend/price)因子
  df['div_factor'] = 1 + df['dividend'] / df['price']
  
  # 计算窗口内div_factor的乘积（窗口大小为n+1=3）
  df['div_product'] = df['div_factor'].rolling(window=3).apply(lambda x: x.prod(), raw=True)
  
  # 计算价格比例项：当前price / 前n行的price
  df['price_ratio'] = df['price'] / df['price'].shift(2)
  
  # 计算最终复利收益，前n行设为NA
  df['return'] = df['price_ratio'] * df['div_product']
  df.loc[:n-1, 'return'] = pd.NA
  
  # 可选：删除中间辅助列
  df = df.drop(['div_factor', 'div_product', 'price_ratio'], axis=1)
  

• 关键说明

  • 使用rolling(window=3)获取包含当前行及前两行的窗口
  • apply(lambda x: x.prod())计算窗口内的乘积，raw=True可提升计算效率
  • shift(2)获取n=2周期前的价格数据
  • 全程采用向量化处理，性能远优于逐行循环

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