""" 动量因子计算模块 支持两种动量因子： 1. N日涨幅（简单动量） 2. 斜率×R²趋势得分（改进版） """ import numpy as np import pandas as pd from sklearn.linear_model import LinearRegression import math def calculate_momentum(price_series: pd.Series, n: int) -> pd.Series: """ 计算 N 日涨幅（简单动量） Args: price_series: 价格序列 n: 动量窗口天数 Returns: Series: N日涨幅 """ return price_series / price_series.shift(n + 1) - 1.0 def _slope_r2_score(srs: pd.Series, n: int = 25) -> float: """ 单次计算斜率×R²趋势得分 Args: srs: 价格窗口序列（长度为 n） n: 窗口长度 Returns: float: 斜率 × R² × 10000 """ if srs.shape[0] < n: return np.nan x = np.arange(1, n + 1).reshape(-1, 1) y = srs.values / srs.values[0] # 归一化 lr = LinearRegression().fit(x, y) slope = lr.coef_[0] r_squared = lr.score(x, y) score = 10000 * slope * r_squared return score def calculate_weighted_momentum_score(prices: np.ndarray) -> float: """ 加权线性回归动量得分 (匹配 动量.py / JoinQuant 逻辑) Args: prices: 价格数组 Returns: float: 年化收益率 * R² """ if len(prices) < 5: return 0.0 y = np.log(prices) x = np.arange(len(y)) weights = np.linspace(1, 2, len(y)) # 近期权重更高 (1 -> 2) # 加权线性回归 # 使用 np.polyfit 的 w 参数进行加权 slope, intercept = np.polyfit(x, y, 1, w=weights) annualized_returns = math.exp(slope * 250) - 1 # 加权R² y_pred = slope * x + intercept ss_res = np.sum(weights * (y - y_pred) ** 2) ss_tot = np.sum(weights * (y - np.average(y, weights=weights)) ** 2) r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0 return annualized_returns * r2 def calculate_atr(high: pd.Series, low: pd.Series, close: pd.Series, period: int) -> pd.Series: """计算ATR（不依赖talib）""" prev_close = close.shift(1) tr = pd.concat([ high - low, (high - prev_close).abs(), (low - prev_close).abs(), ], axis=1).max(axis=1) return tr.rolling(window=period, min_periods=period).mean() def apply_crash_filter(prices: np.ndarray, score: float) -> float: """崩盘过滤：连续3天有任一天跌>5%""" if len(prices) < 4: return score r1 = prices[-1] / prices[-2] r2 = prices[-2] / prices[-3] r3 = prices[-3] / prices[-4] # 条件1：任一天跌>5% con1 = min(r1, r2, r3) < 0.95 # 条件2：连续下跌且累计跌>5% con2 = (r1 < 1) and (r2 < 1) and (r3 < 1) and (prices[-1] / prices[-4] < 0.95) if con1 or con2: return 0.0 return score def calculate_slope_r2(price_series: pd.Series, n: int = 25) -> pd.Series: """ 计算斜率×R²趋势得分序列 Args: price_series: 价格序列 n: 滚动窗口天数 Returns: Series: 趋势得分序列 """ return price_series.rolling(n).apply( lambda x: _slope_r2_score(x, n), raw=False ) def calculate_daily_return(price_series: pd.Series) -> pd.Series: """ 计算日收益率 Args: price_series: 价格序列 Returns: Series: 日收益率 """ return price_series / price_series.shift(1) - 1 def _is_china_index(code: str) -> bool: """判断是否为A股指数""" return code.endswith('.SH') or code.endswith('.SZ') or code.endswith('.SS') def compute_factors( index_data: pd.DataFrame, code_list: list, n: int = 25, factor_type: str = "slope_r2", etf_data: pd.DataFrame = None, code_config: dict = None, index_ohlcv_data: dict = None, auto_day: bool = False, min_days: int = 20, max_days: int = 60, ) -> tuple[pd.DataFrame, list]: """ 计算所有指数的因子和日收益率（横截面策略版本） Args: index_data: 宽格式指数收盘价数据 (对齐后) code_list: 标的代码列表 n: 默认窗口天数 factor_type: 因子类型 ('momentum', 'slope_r2', 'weighted_momentum') etf_data: 宽格式ETF收盘价数据 (用于收益计算) code_config: 代码配置字典 index_ohlcv_data: 原始指数OHLCV数据字典 {code: df} auto_day: 是否启用动态ATR周期 min_days: 动态周期最小值 max_days: 动态周期最大值 """ a_share_dates = index_data.index # 为每个标的单独计算指标，然后对齐到A股交易日历 result = pd.DataFrame(index=a_share_dates) # 使用一个新的列表来存储真正的有效代码 processed_codes = [] for code in code_list: # 优先使用 OHLCV 数据（如果提供） if index_ohlcv_data and code in index_ohlcv_data: df = index_ohlcv_data[code].dropna() else: # 退而求其次使用 index_data 中的 close if code not in index_data: continue df = pd.DataFrame({'close': index_data[code].dropna()}) if len(df) < n + 1: print(f" ⚠ 剔除 {code}: 数据不足 ({len(df)} < {n+1})") continue # 按照该标的自己的交易日历计算指标 if auto_day and 'high' in df.columns and 'low' in df.columns: # 动态周期逻辑 long_atr = calculate_atr(df['high'], df['low'], df['close'], max_days) short_atr = calculate_atr(df['high'], df['low'], df['close'], min_days) # 计算滚动窗口大小 def get_dynamic_n(row, la_col, sa_col): la = row[la_col] sa = row[sa_col] if la > 0 and not np.isnan(la) and not np.isnan(sa): ratio = min(0.9, sa / la) return int(min_days + (max_days - min_days) * (1 - ratio)) return n # 合并ATR到主DF以进行滚动应用 df_temp = df.copy() df_temp['la'] = long_atr df_temp['sa'] = short_atr # 逐日计算得分 (较慢但准确) scores = [] for i in range(len(df_temp)): row = df_temp.iloc[i] d_n = get_dynamic_n(row, 'la', 'sa') if i < d_n: scores.append(np.nan) continue window_prices = df_temp['close'].iloc[i-d_n+1 : i+1].values if factor_type == "weighted_momentum": s = calculate_weighted_momentum_score(window_prices) else: s = _slope_r2_score(pd.Series(window_prices), d_n) # 应用崩盘过滤 s = apply_crash_filter(df_temp['close'].iloc[:i+1].values, s) scores.append(s) factor_series = pd.Series(scores, index=df.index) else: # 固定周期逻辑 if factor_type == "momentum": factor_series = calculate_momentum(df['close'], n) elif factor_type == "slope_r2": factor_series = calculate_slope_r2(df['close'], n) elif factor_type == "weighted_momentum": factor_series = df['close'].rolling(n).apply( lambda x: apply_crash_filter(df['close'].loc[:x.index[-1]].values, calculate_weighted_momentum_score(x.values)), raw=False ) else: raise ValueError(f"不支持的因子类型: {factor_type}") # 对齐到A股交易日历 price_aligned = df['close'].reindex(a_share_dates, method='ffill') factor_aligned = factor_series.reindex(a_share_dates, method='ffill') # 使用传入的ETF数据计算收益（如果有） if etf_data is not None and code in etf_data: return_aligned = calculate_daily_return(etf_data[code].reindex(a_share_dates, method='ffill')) else: return_aligned = calculate_daily_return(price_aligned) result[code] = price_aligned result[f"得分_{code}"] = factor_aligned result[f"日收益率_{code}"] = return_aligned processed_codes.append(code) # 过滤掉缺失值过多的指数 total_rows = len(result) final_valid_codes = [] for code in processed_codes: null_pct = result[code].isnull().sum() / total_rows if null_pct > 0.5: print(f" ⚠ 剔除 {code}: 对齐后缺失率 {null_pct:.1%} 过高") result = result.drop(columns=[code, f"得分_{code}", f"日收益率_{code}"], errors='ignore') else: final_valid_codes.append(code) return result, final_valid_codes