115 lines
5.9 KiB
Python
115 lines
5.9 KiB
Python
import pandas as pd
|
||
import numpy as np
|
||
|
||
def holder_trade_factors(all_data_df: pd.DataFrame, stk_holdertrade_df: pd.DataFrame) -> pd.DataFrame:
|
||
"""
|
||
生成合并的股东增减持因子以及 change_ratio 相关因子（优化版）。
|
||
|
||
Args:
|
||
stk_holdertrade_df (pd.DataFrame): 股东增减持数据，包含 'ts_code', 'ann_date', 'in_de', 'change_ratio'。
|
||
all_data_df (pd.DataFrame): 所有日期所有股票数据，包含 'ts_code', 'trade_date'。
|
||
|
||
Returns:
|
||
pd.DataFrame: 包含增减持因子的 all_data_df。
|
||
"""
|
||
print('正在计算股东增减持因子（优化版）...')
|
||
|
||
# 1. 确保日期列为 datetime 类型
|
||
stk_holdertrade_df['ann_date'] = pd.to_datetime(stk_holdertrade_df['ann_date'])
|
||
all_data_df['trade_date'] = pd.to_datetime(all_data_df['trade_date'])
|
||
|
||
# 2. 对增减持数据进行预处理和排序（排序在此阶段可能不是严格必需的，但保持良好习惯）
|
||
holder_data_processed = stk_holdertrade_df.copy()
|
||
holder_data_processed['change_ratio_in_agg'] = holder_data_processed['change_ratio'].where(holder_data_processed['in_de'] == 'IN', 0)
|
||
holder_data_processed['change_ratio_de_agg'] = holder_data_processed['change_ratio'].where(holder_data_processed['in_de'] == 'DE', 0)
|
||
holder_data_processed['change_ratio_total_agg'] = holder_data_processed['change_ratio']
|
||
holder_data_processed['in_de_numeric'] = holder_data_processed['in_de'].map({'IN': 1, 'DE': -1}).fillna(0) # 用于判断类型
|
||
|
||
# 提前获取所有唯一的交易日期集合，以提高查找效率
|
||
all_trade_dates_set = set(all_data_df['trade_date'].unique())
|
||
|
||
# 3. 构建一个辅助DataFrame，记录每个公告在未来10个日历日（且是交易日）的影响
|
||
expanded_holder_events = []
|
||
for _, row in holder_data_processed.iterrows():
|
||
ts_code = row['ts_code']
|
||
ann_date = row['ann_date']
|
||
|
||
# 生成从公告日期开始的未来10个日历日的日期范围（包括公告日本身）
|
||
# pd.Timedelta(days=10) 表示从公告日+10天
|
||
# pd.date_range(start=ann_date, end=ann_date + pd.Timedelta(days=10), freq='D')
|
||
# 更精确地生成11个日期，涵盖公告日及其后的10个日历日
|
||
future_dates = pd.date_range(start=ann_date, periods=11, freq='D')
|
||
|
||
for date_in_window in future_dates:
|
||
# 只有当日期是实际交易日时才添加
|
||
if date_in_window in all_trade_dates_set:
|
||
expanded_holder_events.append({
|
||
'ts_code': ts_code,
|
||
'trade_date': date_in_window,
|
||
'in_de_numeric': row['in_de_numeric'],
|
||
'change_ratio_total_agg': row['change_ratio_total_agg'],
|
||
'change_ratio_in_agg': row['change_ratio_in_agg'],
|
||
'change_ratio_de_agg': row['change_ratio_de_agg']
|
||
})
|
||
|
||
if not expanded_holder_events: # 如果没有事件，直接返回原始 df
|
||
# 确保返回的DataFrame与原始df具有相同的列和顺序
|
||
# 并填充为默认值
|
||
default_factors = pd.DataFrame({
|
||
'holder_trade_type_10d': None,
|
||
'holder_change_ratio_sum_10d': 0.0,
|
||
'holder_in_change_ratio_sum_10d': 0.0,
|
||
'holder_de_change_ratio_sum_10d': 0.0
|
||
}, index=all_data_df.index)
|
||
return pd.concat([all_data_df, default_factors], axis=1)
|
||
|
||
|
||
expanded_holder_events_df = pd.DataFrame(expanded_holder_events)
|
||
|
||
# 4. 聚合每个 (ts_code, trade_date) 对上的事件
|
||
# 可能会有重复的 (ts_code, trade_date) 对，因为一个交易日可能受多个公告影响
|
||
daily_aggregated_factors = expanded_holder_events_df.groupby(['ts_code', 'trade_date']).agg(
|
||
holder_change_ratio_sum_10d=('change_ratio_total_agg', 'sum'),
|
||
holder_in_change_ratio_sum_10d=('change_ratio_in_agg', 'sum'),
|
||
holder_de_change_ratio_sum_10d=('change_ratio_de_agg', 'sum'),
|
||
# 对于 holder_trade_type_10d，聚合 in_de_numeric 的唯一值集合
|
||
_in_de_types_unique=('in_de_numeric', lambda x: set(x)) # 获取该日期窗口内所有独特的增减持类型
|
||
).reset_index()
|
||
|
||
# 根据 _in_de_types_unique 确定 holder_trade_type_10d
|
||
def get_trade_type(unique_types_set):
|
||
if 1 in unique_types_set and -1 in unique_types_set:
|
||
return 'BOTH'
|
||
elif 1 in unique_types_set:
|
||
return 'IN'
|
||
elif -1 in unique_types_set:
|
||
return 'DE'
|
||
else:
|
||
return None # 理论上不应该发生，除非 unique_types_set 为空或只包含0
|
||
|
||
daily_aggregated_factors['holder_trade_type_10d'] = daily_aggregated_factors['_in_de_types_unique'].apply(get_trade_type)
|
||
|
||
# 移除辅助列
|
||
daily_aggregated_factors.drop(columns=['_in_de_types_unique'], inplace=True)
|
||
|
||
# 5. 将计算得到的因子合并回 all_data_df
|
||
# 确保 all_data_df 也按 ts_code, trade_date 排序，以便 merge 高效
|
||
all_data_df_sorted = all_data_df.sort_values(['ts_code', 'trade_date']).reset_index(drop=True)
|
||
|
||
final_df = pd.merge(
|
||
all_data_df_sorted,
|
||
daily_aggregated_factors,
|
||
on=['ts_code', 'trade_date'],
|
||
how='left'
|
||
)
|
||
|
||
# 6. 对于没有增减持记录的日期，因子值为 None 或 0
|
||
# 在 merge_asof 中确实需要排序，但在这种事件展开的方法中，merge 是普通的 left merge，不需要预排序。
|
||
# 考虑到最终的 merge，最好还是保持 `all_data_df` 和 `daily_aggregated_factors` 的键排序。
|
||
# 所以在 `merge` 前对 `all_data_df` 进行一次排序是好的实践。
|
||
|
||
final_df['holder_trade_type_10d'] = final_df['holder_trade_type_10d'].fillna(None)
|
||
fillna_ratio_cols = ['holder_change_ratio_sum_10d', 'holder_in_change_ratio_sum_10d', 'holder_de_change_ratio_sum_10d']
|
||
final_df[fillna_ratio_cols] = final_df[fillna_ratio_cols].fillna(0.0)
|
||
|
||
return final_df |