使用Pandas掩码筛选DataFrame并生成含历史/未来收益的新表

原始DataFrame：

date   price   MA1  MA2 MA3
date0  price0  12   10  8
date1  price1  11   11 11 
date2  price2  12   21 14
date3  price3  13   12 15
date4  price4  14   14 14
date5  price5  15   17 14
date6  price6  19   16 15
date7  price7  15   12 13
date8  price8  11   10 13
date9  price9  21   12 13
date10 price10 13   11 14
date11 price11 14   14 14
date12 price12 16   16 16
date13 price13 34   32 23
date14 price14 12   12 12

已用以下代码筛选出MA1、MA2、MA3数值相等的行：

df =(data
    .assign(same=lambda x: (x['MA1'] == x['MA2']) & (x['MA1'] == x['MA3']))
    .loc[lambda x: x.same == True]
)

筛选结果：

date   price   MA1  MA2 MA3
date1  price1  11   11 11 
date4  price4  14   14 14
date11 price11 14   14 14
date12 price12 16   16 16
date14 price14 12   12 12

需要生成包含前一期价格、后一期价格、前一期收益率、后一期收益率的新DataFrame，格式如下：

date   price  price_past price_fut      return_past                return_future            
date1  price1  price0    price2      (price1-price0)/price0     (price2-price1)/price1     
date4  price4  price3    price5      (price4-price3)/price3     (price5-price4)/price4      
date11 price11 price10   price12     (price11-price10)/price10  (price12-price11)/price11  
date12 price12 price11   price13     (price12-price11)/price11  (price13 -price12)/price12                
date14 price14 price13   price15     (price14-price13)/price13  (price15-price14)/price14   

可以利用Pandas的shift()方法获取前后期价格，再推导收益率，具体实现代码如下：

# 1. 执行原筛选逻辑，移除辅助列
filtered_df = (data
    .assign(same=lambda x: (x['MA1'] == x['MA2']) & (x['MA1'] == x['MA3']))
    .loc[lambda x: x.same == True]
    .drop(columns='same')
)

# 2. 从原始数据提取价格序列，用于匹配前后期价格
price_map = data.set_index('date')['price']

# 3. 添加前一期、后一期价格列
filtered_df['price_past'] = filtered_df['date'].map(price_map.shift(1))
filtered_df['price_fut'] = filtered_df['date'].map(price_map.shift(-1))

# 4. 计算收益率
filtered_df['return_past'] = (filtered_df['price'] - filtered_df['price_past']) / filtered_df['price_past']
filtered_df['return_future'] = (filtered_df['price_fut'] - filtered_df['price']) / filtered_df['price']

# 5. 调整列顺序为需求格式
final_df = filtered_df[['date', 'price', 'price_past', 'price_fut', 'return_past', 'return_future']]

关键说明： #

• price_map.shift(1)将价格序列整体下移一位，得到每个日期的前一期价格；shift(-1)上移一位得到后一期价格
• 使用map()方法可以快速将筛选后的日期对应到前后期价格
• 收益率直接按照(当期价格-对比期价格)/对比期价格的公式计算即可
• 如果原始数据的date列是索引，可直接在筛选后的DataFrame上调用shift()，无需单独生成价格映射

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