factor优化，改为polars

main/factor/polars_volatility_factors.py

@@ -0,0 +1,419 @@
+"""
+波动率因子 - 使用Polars实现
+包含上行波动率、下行波动率、波动率比率等相关因子计算
+"""
+
+import polars as pl
+import numpy as np
+from typing import Dict, List, Optional, Any
+from operator_framework import StockWiseOperator, OperatorConfig
+
+
+class UpsideVolatilityOperator(StockWiseOperator):
+    """上行波动率算子"""
+    
+    def __init__(self, window: int = 20):
+        config = OperatorConfig(
+            name=f"upside_volatility_{window}",
+            description=f"{window}日上行波动率",
+            required_columns=['pct_chg'],
+            output_columns=[f'upside_volatility_{window}'],
+            parameters={'window': window}
+        )
+        super().__init__(config)
+        self.window = window
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算上行波动率"""
+        # 分离正收益率
+        pos_returns = pl.when(pl.col('pct_chg') > 0).then(pl.col('pct_chg')).otherwise(0)
+        
+        # 计算正收益率的平方
+        pos_returns_sq = pos_returns.pow(2)
+        
+        # 计算滚动和
+        rolling_pos_count = (pl.col('pct_chg') > 0).rolling_sum(window=self.window)
+        rolling_pos_sum = pos_returns.rolling_sum(window=self.window)
+        rolling_pos_sum_sq = pos_returns_sq.rolling_sum(window=self.window)
+        
+        # 计算方差和标准差
+        pos_mean_sq = rolling_pos_sum_sq / rolling_pos_count
+        pos_mean = rolling_pos_sum / rolling_pos_count
+        pos_var = pos_mean_sq - pos_mean.pow(2)
+        
+        # 处理样本数不足的情况
+        pos_var = pl.when(rolling_pos_count >= 2).then(pos_var).otherwise(None)
+        pos_var = pos_var.clip(lower=0)
+        
+        upside_vol = pos_var.sqrt()
+        
+        return stock_df.with_columns(upside_vol.alias(f'upside_volatility_{self.window}'))
+
+
+class DownsideVolatilityOperator(StockWiseOperator):
+    """下行波动率算子"""
+    
+    def __init__(self, window: int = 20):
+        config = OperatorConfig(
+            name=f"downside_volatility_{window}",
+            description=f"{window}日下行波动率",
+            required_columns=['pct_chg'],
+            output_columns=[f'downside_volatility_{window}'],
+            parameters={'window': window}
+        )
+        super().__init__(config)
+        self.window = window
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算下行波动率"""
+        # 分离负收益率
+        neg_returns = pl.when(pl.col('pct_chg') < 0).then(pl.col('pct_chg')).otherwise(0)
+        
+        # 计算负收益率的平方
+        neg_returns_sq = neg_returns.pow(2)
+        
+        # 计算滚动和
+        rolling_neg_count = (pl.col('pct_chg') < 0).rolling_sum(window=self.window)
+        rolling_neg_sum = neg_returns.rolling_sum(window=self.window)
+        rolling_neg_sum_sq = neg_returns_sq.rolling_sum(window=self.window)
+        
+        # 计算方差和标准差
+        neg_mean_sq = rolling_neg_sum_sq / rolling_neg_count
+        neg_mean = rolling_neg_sum / rolling_neg_count
+        neg_var = neg_mean_sq - neg_mean.pow(2)
+        
+        # 处理样本数不足的情况
+        neg_var = pl.when(rolling_neg_count >= 2).then(neg_var).otherwise(None)
+        neg_var = neg_var.clip(lower=0)
+        
+        downside_vol = neg_var.sqrt()
+        
+        return stock_df.with_columns(downside_vol.alias(f'downside_volatility_{self.window}'))
+
+
+class VolatilityRatioOperator(StockWiseOperator):
+    """波动率比率算子"""
+    
+    def __init__(self, window: int = 20):
+        config = OperatorConfig(
+            name=f"volatility_ratio_{window}",
+            description=f"{window}日波动率比率",
+            required_columns=['pct_chg'],
+            output_columns=[f'volatility_ratio_{window}'],
+            parameters={'window': window}
+        )
+        super().__init__(config)
+        self.window = window
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算波动率比率"""
+        # 计算上行和下行波动率
+        pos_returns = pl.when(pl.col('pct_chg') > 0).then(pl.col('pct_chg')).otherwise(0)
+        neg_returns = pl.when(pl.col('pct_chg') < 0).then(pl.col('pct_chg')).otherwise(0)
+        
+        # 计算滚动标准差
+        upside_vol = pos_returns.rolling_std(window=self.window)
+        downside_vol = neg_returns.rolling_std(window=self.window)
+        
+        # 计算比率
+        vol_ratio = upside_vol / downside_vol
+        
+        # 处理无穷大和NaN值
+        vol_ratio = vol_ratio.replace([np.inf, -np.inf], None).fill_null(0)
+        
+        return stock_df.with_columns(vol_ratio.alias(f'volatility_ratio_{self.window}'))
+
+
+class ReturnSkewnessOperator(StockWiseOperator):
+    """收益率偏度算子"""
+    
+    def __init__(self, window: int = 5):
+        config = OperatorConfig(
+            name=f"return_skewness_{window}",
+            description=f"{window}日收益率偏度",
+            required_columns=['pct_chg'],
+            output_columns=[f'return_skewness_{window}'],
+            parameters={'window': window}
+        )
+        super().__init__(config)
+        self.window = window
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算收益率偏度"""
+        skewness = pl.col('pct_chg').rolling_skew(window=self.window)
+        
+        return stock_df.with_columns(skewness.alias(f'return_skewness_{self.window}'))
+
+
+class ReturnKurtosisOperator(StockWiseOperator):
+    """收益率峰度算子"""
+    
+    def __init__(self, window: int = 5):
+        config = OperatorConfig(
+            name=f"return_kurtosis_{window}",
+            description=f"{window}日收益率峰度",
+            required_columns=['pct_chg'],
+            output_columns=[f'return_kurtosis_{window}'],
+            parameters={'window': window}
+        )
+        super().__init__(config)
+        self.window = window
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算收益率峰度"""
+        kurtosis = pl.col('pct_chg').rolling_kurt(window=self.window)
+        
+        return stock_df.with_columns(kurtosis.alias(f'return_kurtosis_{self.window}'))
+
+
+class VolatilityAmplificationOperator(StockWiseOperator):
+    """亏损状态波动率放大算子"""
+    
+    def __init__(self, n: int = 20):
+        config = OperatorConfig(
+            name=f"vol_amp_loss_{n}",
+            description=f"{n}日亏损状态波动率放大",
+            required_columns=['pct_chg', 'weight_avg', 'close'],
+            output_columns=[f'vol_amp_loss_{n}'],
+            parameters={'n': n}
+        )
+        super().__init__(config)
+        self.n = n
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算亏损状态波动率放大"""
+        # 计算n日波动率
+        vol_n = pl.col('pct_chg').rolling_std(window=self.n)
+        
+        # 计算亏损程度
+        loss_degree = pl.max_horizontal(0, pl.col('weight_avg') - pl.col('close')) / pl.col('close')
+        
+        # 计算放大因子
+        vol_amp = vol_n * loss_degree
+        
+        return stock_df.with_columns(vol_amp.alias(f'vol_amp_loss_{self.n}'))
+
+
+class HighVolDropWhenProfitableOperator(StockWiseOperator):
+    """高成交量下跌当获利状态算子"""
+    
+    def __init__(self, n: int = 20, m: int = 5, profit_thresh: float = 0.1, 
+                 drop_thresh: float = -0.03, vol_multiple: float = 2.0):
+        config = OperatorConfig(
+            name=f"vol_drop_profit_cnt_{m}",
+            description=f"{m}日高成交量下跌当获利状态计数",
+            required_columns=['close', 'pct_chg', 'vol', 'weight_avg'],
+            output_columns=[f'vol_drop_profit_cnt_{m}'],
+            parameters={'n': n, 'm': m, 'profit_thresh': profit_thresh, 
+                       'drop_thresh': drop_thresh, 'vol_multiple': vol_multiple}
+        )
+        super().__init__(config)
+        self.n = n
+        self.m = m
+        self.profit_thresh = profit_thresh
+        self.drop_thresh = drop_thresh
+        self.vol_multiple = vol_multiple
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算高成交量下跌当获利状态计数"""
+        # 判断是否获利
+        is_profitable = pl.col('close') > pl.col('weight_avg') * (1 + self.profit_thresh)
+        
+        # 判断是否下跌
+        is_dropping = pl.col('pct_chg') < self.drop_thresh
+        
+        # 计算滚动均值和标准差
+        rolling_mean_vol = pl.col('vol').rolling_mean(window=self.n)
+        rolling_std_vol = pl.col('vol').rolling_std(window=self.n).fill_null(0)
+        
+        # 判断是否高成交量
+        is_high_vol = pl.col('vol') > (rolling_mean_vol + self.vol_multiple * rolling_std_vol)
+        
+        # 计算事件
+        event = is_profitable & is_dropping & is_high_vol
+        
+        # 计算m日累计
+        event_cnt = event.cast(int).rolling_sum(window=self.m)
+        
+        return stock_df.with_columns(event_cnt.alias(f'vol_drop_profit_cnt_{self.m}'))
+
+
+class LargeFlowVolatilityInteractionOperator(StockWiseOperator):
+    """大单资金流驱动波动率交互算子"""
+    
+    def __init__(self, n: int = 20):
+        config = OperatorConfig(
+            name=f"lg_flow_vol_interact_{n}",
+            description=f"{n}日大单资金流驱动波动率交互",
+            required_columns=['pct_chg', 'buy_lg_vol', 'buy_elg_vol', 'sell_lg_vol', 
+                            'sell_elg_vol', 'vol', 'close'],
+            output_columns=[f'lg_flow_vol_interact_{n}'],
+            parameters={'n': n}
+        )
+        super().__init__(config)
+        self.n = n
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算大单资金流驱动波动率交互"""
+        epsilon = 1e-8
+        
+        # 计算n日波动率
+        vol_n = pl.col('pct_chg').rolling_std(window=self.n)
+        
+        # 计算大单净额
+        net_lg_flow_val = (
+            (pl.col('buy_lg_vol') + pl.col('buy_elg_vol') - 
+             pl.col('sell_lg_vol') - pl.col('sell_elg_vol')) * pl.col('close')
+        )
+        
+        # 计算总成交额
+        total_val = pl.col('vol') * pl.col('close')
+        
+        # 计算大单净流入比率绝对值
+        abs_net_lg_flow_ratio = net_lg_flow_val.abs() / (total_val + epsilon)
+        
+        # 计算n日均值
+        abs_ratio_n = abs_net_lg_flow_ratio.rolling_mean(window=self.n)
+        
+        # 计算交互项
+        interaction = vol_n * abs_ratio_n
+        
+        return stock_df.with_columns(interaction.alias(f'lg_flow_vol_interact_{self.n}'))
+
+
+class VolatilityAdjustedROCPOperator(StockWiseOperator):
+    """波动率调整收益率算子"""
+    
+    def __init__(self, n: int = 20):
+        config = OperatorConfig(
+            name=f"vol_adj_roc_{n}",
+            description=f"{n}日波动率调整收益率",
+            required_columns=['close', 'pct_chg'],
+            output_columns=[f'vol_adj_roc_{n}'],
+            parameters={'n': n}
+        )
+        super().__init__(config)
+        self.n = n
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算波动率调整收益率"""
+        # 计算n日收益率
+        roc_n = pl.col('close').pct_change(self.n)
+        
+        # 计算n日波动率
+        vol_n = pl.col('pct_chg').rolling_std(window=self.n).fill_null(0)
+        
+        # 计算波动率调整收益率
+        vol_adj_roc = roc_n / (vol_n + 1e-10)  # 避免除零
+        
+        return stock_df.with_columns(vol_adj_roc.alias(f'vol_adj_roc_{self.n}'))
+
+
+class StandardDeviation5Operator(StockWiseOperator):
+    """5日收益率标准差算子"""
+    
+    def __init__(self):
+        config = OperatorConfig(
+            name="std_return_5",
+            description="5日收益率标准差",
+            required_columns=['close'],
+            output_columns=['std_return_5'],
+            parameters={}
+        )
+        super().__init__(config)
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算5日收益率标准差"""
+        # 计算收益率
+        returns = pl.col('close').pct_change()
+        
+        # 计算5日标准差
+        std_5 = returns.rolling_std(window=5)
+        
+        return stock_df.with_columns(std_5.alias('std_return_5'))
+
+
+class StandardDeviation90Operator(StockWiseOperator):
+    """90日收益率标准差算子"""
+    
+    def __init__(self):
+        config = OperatorConfig(
+            name="std_return_90",
+            description="90日收益率标准差",
+            required_columns=['close'],
+            output_columns=['std_return_90'],
+            parameters={}
+        )
+        super().__init__(config)
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算90日收益率标准差"""
+        # 计算收益率
+        returns = pl.col('close').pct_change()
+        
+        # 计算90日标准差
+        std_90 = returns.rolling_std(window=90)
+        
+        return stock_df.with_columns(std_90.alias('std_return_90'))
+
+
+class StandardDeviation90ShiftedOperator(StockWiseOperator):
+    """90日收益率标准差(移位)算子"""
+    
+    def __init__(self):
+        config = OperatorConfig(
+            name="std_return_90_2",
+            description="90日收益率标准差(移位10日)",
+            required_columns=['close'],
+            output_columns=['std_return_90_2'],
+            parameters={}
+        )
+        super().__init__(config)
+    
+    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
+        """计算90日收益率标准差(移位10日)"""
+        # 计算收益率(移位10日)
+        returns = pl.col('close').shift(10).pct_change()
+        
+        # 计算90日标准差
+        std_90_2 = returns.rolling_std(window=90)
+        
+        return stock_df.with_columns(std_90_2.alias('std_return_90_2'))
+
+
+# 波动率因子集合
+VOLATILITY_OPERATORS = [
+    UpsideVolatilityOperator(),
+    DownsideVolatilityOperator(),
+    VolatilityRatioOperator(),
+    ReturnSkewnessOperator(),
+    ReturnKurtosisOperator(),
+    VolatilityAmplificationOperator(),
+    HighVolDropWhenProfitableOperator(),
+    LargeFlowVolatilityInteractionOperator(),
+    VolatilityAdjustedROCPOperator(),
+    StandardDeviation5Operator(),
+    StandardDeviation90Operator(),
+    StandardDeviation90ShiftedOperator(),
+]
+
+
+def apply_volatility_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
+    """
+    应用所有波动率因子
+    
+    Args:
+        df: 输入的Polars DataFrame
+        operators: 要应用的算子列表，如果为None则使用默认列表
+    
+    Returns:
+        添加了波动率因子的DataFrame
+    """
+    if operators is None:
+        operators = VOLATILITY_OPERATORS
+    
+    result_df = df
+    for operator in operators:
+        result_df = operator(result_df)
+    
+    return result_df