411 lines
15 KiB
Python
411 lines
15 KiB
Python
"""
|
||
动量因子 - 使用Polars实现
|
||
包含动量、趋势、均线等相关因子计算
|
||
"""
|
||
|
||
import polars as pl
|
||
import numpy as np
|
||
from typing import Dict, List, Optional, Any
|
||
|
||
from tqdm import tqdm
|
||
from main.factor.operator_framework import StockWiseOperator, OperatorConfig
|
||
from scipy.stats import linregress
|
||
|
||
|
||
class PriceMinusDeductionPriceOperator(StockWiseOperator):
|
||
"""价格减抵扣价算子"""
|
||
|
||
def __init__(self, n: int = 10):
|
||
if n <= 0:
|
||
raise ValueError("n must be positive")
|
||
config = OperatorConfig(
|
||
name=f"price_minus_deduction_price_{n}",
|
||
description=f"{n}日价格减抵扣价",
|
||
required_columns=['close'],
|
||
output_columns=[f'price_minus_deduction_price_{n}'],
|
||
parameters={'n': n}
|
||
)
|
||
super().__init__(config)
|
||
self.n = n
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'price_minus_deduction_price_{self.n}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
# 抵扣价是 n 日前的价格（更合理），若坚持 n-1 则保留
|
||
deduction_price = group_df['close'].shift(self.n) # 建议用 n，不是 n-1
|
||
price_diff = group_df['close'] - deduction_price
|
||
return price_diff.alias(self.get_factor_name())
|
||
|
||
|
||
class PriceDeductionPriceDiffRatioToSMAOperator(StockWiseOperator):
|
||
"""价格抵扣价差值相对SMA比率算子"""
|
||
|
||
def __init__(self, n: int = 10):
|
||
if n <= 0:
|
||
raise ValueError("n must be positive")
|
||
config = OperatorConfig(
|
||
name=f"price_deduction_price_diff_ratio_to_sma_{n}",
|
||
description=f"{n}日价格抵扣价差值相对SMA比率",
|
||
required_columns=['close'],
|
||
output_columns=[f'price_deduction_price_diff_ratio_to_sma_{n}'],
|
||
parameters={'n': n}
|
||
)
|
||
super().__init__(config)
|
||
self.n = n
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'price_deduction_price_diff_ratio_to_sma_{self.n}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
sma = group_df['close'].rolling_mean(window_size=self.n)
|
||
deduction_price = group_df['close'].shift(self.n)
|
||
diff = group_df['close'] - deduction_price
|
||
ratio = diff / (sma + 1e-8)
|
||
return ratio.alias(self.get_factor_name())
|
||
|
||
|
||
class CatPriceVsSmaVsDeductionPriceOperator(StockWiseOperator):
|
||
"""价格vsSMAvs抵扣价分类算子"""
|
||
|
||
def __init__(self, n: int = 10):
|
||
if n <= 0:
|
||
raise ValueError("n must be positive")
|
||
config = OperatorConfig(
|
||
name=f"cat_price_vs_sma_vs_deduction_price_{n}",
|
||
description=f"{n}日价格vsSMAvs抵扣价分类",
|
||
required_columns=['close'],
|
||
output_columns=[f'cat_price_vs_sma_vs_deduction_price_{n}'],
|
||
parameters={'n': n}
|
||
)
|
||
super().__init__(config)
|
||
self.n = n
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'cat_price_vs_sma_vs_deduction_price_{self.n}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
sma = group_df['close'].rolling_mean(window_size=self.n)
|
||
deduction_price = group_df['close'].shift(self.n)
|
||
|
||
cond1 = (group_df['close'] > sma) & (deduction_price > sma)
|
||
cond2 = (group_df['close'] < sma) & (deduction_price < sma)
|
||
cond3 = (group_df['close'] > sma) & (deduction_price <= sma)
|
||
cond4 = (group_df['close'] <= sma) & (deduction_price > sma)
|
||
|
||
classification = (
|
||
pl.when(cond1).then(1)
|
||
.when(cond2).then(2)
|
||
.when(cond3).then(3)
|
||
.when(cond4).then(4)
|
||
.otherwise(0)
|
||
)
|
||
return classification.alias(self.get_factor_name())
|
||
|
||
|
||
# ✅ 修复：使用 rolling_map
|
||
class VolatilitySlopeOperator(StockWiseOperator):
|
||
"""波动率斜率算子"""
|
||
|
||
def __init__(self, long_window: int = 20, short_window: int = 5):
|
||
if long_window <= 0 or short_window <= 0:
|
||
raise ValueError("Windows must be positive")
|
||
config = OperatorConfig(
|
||
name=f"volatility_slope_{long_window}_{short_window}",
|
||
description=f"{long_window}日波动率{short_window}日斜率",
|
||
required_columns=['pct_chg'],
|
||
output_columns=[f'volatility_slope_{long_window}_{short_window}'],
|
||
parameters={'long_window': long_window, 'short_window': short_window}
|
||
)
|
||
super().__init__(config)
|
||
self.long_window = long_window
|
||
self.short_window = short_window
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'volatility_slope_{self.long_window}_{self.short_window}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
# 先计算长期波动率（标准差）
|
||
long_vol = group_df['pct_chg'].rolling_std(window_size=self.long_window)
|
||
|
||
# 定义斜率函数（输入是 numpy array）
|
||
def slope_func(window_vals: np.ndarray) -> float:
|
||
if len(window_vals) < 2 or pl.Series(window_vals).is_null().any():
|
||
return 0.0
|
||
x = np.arange(len(window_vals))
|
||
try:
|
||
slope, _, _, _, _ = linregress(x, window_vals)
|
||
return slope if np.isfinite(slope) else 0.0
|
||
except:
|
||
return 0.0
|
||
|
||
# 对波动率序列应用 rolling_map
|
||
volatility_slope = long_vol.rolling_map(
|
||
function=slope_func,
|
||
window_size=self.short_window,
|
||
min_periods=2 # 至少2点才能算斜率
|
||
)
|
||
return volatility_slope.alias(self.get_factor_name())
|
||
|
||
|
||
# ✅ 修复：使用 rolling_map
|
||
class TurnoverRateTrendStrengthOperator(StockWiseOperator):
|
||
"""换手率趋势强度算子"""
|
||
|
||
def __init__(self, window: int = 5):
|
||
if window <= 0:
|
||
raise ValueError("Window must be positive")
|
||
config = OperatorConfig(
|
||
name=f"turnover_trend_strength_{window}",
|
||
description=f"{window}日换手率趋势强度",
|
||
required_columns=['turnover_rate'],
|
||
output_columns=[f'turnover_trend_strength_{window}'],
|
||
parameters={'window': window}
|
||
)
|
||
super().__init__(config)
|
||
self.window = window
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'turnover_trend_strength_{self.window}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
def slope_func(window_vals: np.ndarray) -> float:
|
||
if len(window_vals) < 2 or pl.Series(window_vals).is_null().any():
|
||
return 0.0
|
||
x = np.arange(len(window_vals))
|
||
try:
|
||
slope, _, _, _, _ = linregress(x, window_vals)
|
||
return slope if np.isfinite(slope) else 0.0
|
||
except:
|
||
return 0.0
|
||
|
||
trend_strength = group_df['turnover_rate'].rolling_map(
|
||
function=slope_func,
|
||
window_size=self.window,
|
||
min_periods=2
|
||
)
|
||
return trend_strength.alias(self.get_factor_name())
|
||
|
||
|
||
class FreeFloatTurnoverSurgeOperator(StockWiseOperator):
|
||
"""自由流通股换手率激增算子"""
|
||
|
||
def __init__(self, window: int = 10):
|
||
if window <= 0:
|
||
raise ValueError("Window must be positive")
|
||
config = OperatorConfig(
|
||
name=f"ff_turnover_surge_{window}",
|
||
description=f"{window}日自由流通股换手率激增",
|
||
required_columns=['turnover_rate'],
|
||
output_columns=[f'ff_turnover_surge_{window}'],
|
||
parameters={'window': window}
|
||
)
|
||
super().__init__(config)
|
||
self.window = window
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'ff_turnover_surge_{self.window}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
avg_turnover = group_df['turnover_rate'].rolling_mean(window_size=self.window)
|
||
surge_ratio = group_df['turnover_rate'] / (avg_turnover + 1e-8)
|
||
return surge_ratio.alias(self.get_factor_name())
|
||
|
||
|
||
class PriceVolumeTrendCoherenceOperator(StockWiseOperator):
|
||
"""价量趋势一致性算子"""
|
||
|
||
def __init__(self, price_window: int = 5, volume_window: int = 20):
|
||
if price_window <= 0 or volume_window <= 0:
|
||
raise ValueError("Windows must be positive")
|
||
config = OperatorConfig(
|
||
name=f"price_volume_coherence_{price_window}_{volume_window}",
|
||
description=f"{price_window}日价格{volume_window}日成交量趋势一致性",
|
||
required_columns=['close', 'vol'],
|
||
output_columns=[f'price_volume_coherence_{price_window}_{volume_window}'],
|
||
parameters={'price_window': price_window, 'volume_window': volume_window}
|
||
)
|
||
super().__init__(config)
|
||
self.price_window = price_window
|
||
self.volume_window = volume_window
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'price_volume_coherence_{self.price_window}_{self.volume_window}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
price_up = (group_df['close'].diff() > 0).cast(pl.Int8)
|
||
price_up_ratio = price_up.rolling_mean(window_size=self.price_window)
|
||
|
||
vol_avg = group_df['vol'].rolling_mean(window_size=self.volume_window)
|
||
vol_above = (group_df['vol'] > vol_avg).cast(pl.Int8)
|
||
vol_above_ratio = vol_above.rolling_mean(window_size=self.price_window)
|
||
|
||
coherence = price_up_ratio * vol_above_ratio
|
||
return coherence.alias(self.get_factor_name())
|
||
|
||
|
||
class FreeFloatToTotalTurnoverRatioOperator(StockWiseOperator):
|
||
"""自由流通股对总换手率比率算子"""
|
||
|
||
def __init__(self):
|
||
config = OperatorConfig(
|
||
name="ff_to_total_turnover_ratio",
|
||
description="自由流通股对总换手率比率",
|
||
required_columns=['turnover_rate'],
|
||
output_columns=['ff_to_total_turnover_ratio'],
|
||
parameters={}
|
||
)
|
||
super().__init__(config)
|
||
|
||
def get_factor_name(self) -> str:
|
||
return 'ff_to_total_turnover_ratio'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
# 实际业务中可能需要 total_turnover_rate，这里简化
|
||
ratio = pl.lit(1.0) # 或根据实际逻辑修改
|
||
return ratio.alias('ff_to_total_turnover_ratio')
|
||
|
||
|
||
class VarianceOperator(StockWiseOperator):
|
||
"""方差算子"""
|
||
|
||
def __init__(self, window: int):
|
||
if window <= 0:
|
||
raise ValueError("Window must be positive")
|
||
config = OperatorConfig(
|
||
name=f"variance_{window}",
|
||
description=f"{window}日方差",
|
||
required_columns=['pct_chg'],
|
||
output_columns=[f'variance_{window}'],
|
||
parameters={'window': window}
|
||
)
|
||
super().__init__(config)
|
||
self.window = window
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'variance_{self.window}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
variance = group_df['pct_chg'].rolling_var(window_size=self.window)
|
||
return variance.alias(self.get_factor_name())
|
||
|
||
|
||
class LimitUpDownOperator(StockWiseOperator):
|
||
"""涨跌停算子"""
|
||
|
||
def __init__(self):
|
||
config = OperatorConfig(
|
||
name="limit_up_down",
|
||
description="涨跌停因子",
|
||
required_columns=['close', 'up_limit', 'down_limit'],
|
||
output_columns=['cat_up_limit', 'cat_down_limit', 'up_limit_count_10d', 'down_limit_count_10d'],
|
||
parameters={}
|
||
)
|
||
super().__init__(config)
|
||
|
||
def get_factor_name(self) -> str:
|
||
return 'cat_up_limit'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
up_limit = (group_df['close'] == group_df['up_limit']).cast(pl.Int8)
|
||
return up_limit.alias('cat_up_limit')
|
||
|
||
|
||
class ConsecutiveUpLimitOperator(StockWiseOperator):
|
||
"""连续涨停算子"""
|
||
|
||
def __init__(self):
|
||
config = OperatorConfig(
|
||
name="consecutive_up_limit",
|
||
description="连续涨停天数",
|
||
required_columns=['cat_up_limit'],
|
||
output_columns=['consecutive_up_limit'],
|
||
parameters={}
|
||
)
|
||
super().__init__(config)
|
||
|
||
def get_factor_name(self) -> str:
|
||
return 'consecutive_up_limit'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
# 简化版：实际连续计数需用 cumsum + groupby trick
|
||
# 这里先返回原始值，后续可优化
|
||
return group_df['cat_up_limit'].alias('consecutive_up_limit')
|
||
|
||
|
||
class MomentumFactorOperator(StockWiseOperator):
|
||
"""动量因子算子"""
|
||
|
||
def __init__(self, alpha: float = 0.5):
|
||
if not (0 <= alpha <= 1):
|
||
raise ValueError("alpha should be between 0 and 1")
|
||
config = OperatorConfig(
|
||
name=f"momentum_factor_{alpha}",
|
||
description=f"动量因子(alpha={alpha})",
|
||
required_columns=['volume_change_rate', 'turnover_deviation'],
|
||
output_columns=[f'momentum_factor_{alpha}'],
|
||
parameters={'alpha': alpha}
|
||
)
|
||
super().__init__(config)
|
||
self.alpha = alpha
|
||
|
||
def get_factor_name(self) -> str:
|
||
return f'momentum_factor_{self.alpha}'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
momentum = group_df['volume_change_rate'] + self.alpha * group_df['turnover_deviation']
|
||
return momentum.alias(self.get_factor_name())
|
||
|
||
|
||
class ResonanceFactorOperator(StockWiseOperator):
|
||
"""共振因子算子"""
|
||
|
||
def __init__(self):
|
||
config = OperatorConfig(
|
||
name="resonance_factor",
|
||
description="共振因子",
|
||
required_columns=['volume_ratio', 'pct_chg'],
|
||
output_columns=['resonance_factor'],
|
||
parameters={}
|
||
)
|
||
super().__init__(config)
|
||
|
||
def get_factor_name(self) -> str:
|
||
return 'resonance_factor'
|
||
|
||
def calc_factor(self, group_df: pl.DataFrame, **kwargs) -> pl.Series:
|
||
resonance = group_df['volume_ratio'] * group_df['pct_chg']
|
||
return resonance.alias('resonance_factor')
|
||
|
||
|
||
# 动量因子集合
|
||
MOMENTUM_OPERATORS = [
|
||
PriceMinusDeductionPriceOperator(10),
|
||
PriceDeductionPriceDiffRatioToSMAOperator(10),
|
||
CatPriceVsSmaVsDeductionPriceOperator(10),
|
||
# VolatilitySlopeOperator(20, 5),
|
||
# TurnoverRateTrendStrengthOperator(5),
|
||
FreeFloatTurnoverSurgeOperator(10),
|
||
PriceVolumeTrendCoherenceOperator(5, 20),
|
||
FreeFloatToTotalTurnoverRatioOperator(),
|
||
VarianceOperator(20),
|
||
LimitUpDownOperator(),
|
||
ConsecutiveUpLimitOperator(),
|
||
# MomentumFactorOperator(0.5),
|
||
ResonanceFactorOperator(),
|
||
]
|
||
|
||
|
||
def apply_momentum_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
|
||
"""
|
||
应用所有动量因子
|
||
"""
|
||
if operators is None:
|
||
operators = MOMENTUM_OPERATORS
|
||
|
||
result_df = df
|
||
for operator in tqdm(operators, desc="Applying momentum factors"):
|
||
result_df = operator.apply(result_df)
|
||
|
||
return result_df
|