Qionk/a-share-quant

1"""2统一数据预处理管线3- 生成日收益率（所有模型的预测目标）4- 生成成交量衍生特征5- 时间序列划分（禁止随机）6"""7 8import numpy as np9import pandas as pd10 11 12def preprocess_data(df: pd.DataFrame, return_clip: float = 0.10,13                     forecast_days: int = 1,14                     index_returns: pd.Series = None) -> pd.DataFrame:15    """16    统一数据预处理：计算日收益率目标 + 成交量衍生特征17 18    返回的 DataFrame 新增列：19      日收益率, future_ret, 目标收益率, 涨跌标签, 目标涨跌,20      成交量变化率, 相对成交量, 量价配合度, 放量上涨, 缩量下跌21 22    df 必须已有列: close, volume, pct_change（或从close计算）23    forecast_days: 持有天数，N日持有期收益 = close.pct_change(N).shift(-N)24    index_returns: 可选，板块/大盘日收益率Series（index对齐），用于计算相对强弱25    """26    df = df.copy()27 28    # ── 1. 日收益率（小数形式，如 0.015 = 1.5%） ──29    if 'pct_change' in df.columns:30        # AKShare/Tencent 返回的 pct_change 是百分比形式（1.5 代表 1.5%），转为小数31        df['日收益率'] = df['pct_change'] / 100.032    else:33        df['日收益率'] = df['close'].pct_change()  # 小数形式，如 0.015 = 1.5%34 35    # 截断 ±10%（符合A股涨跌幅限制）36    df['日收益率'] = df['日收益率'].clip(-return_clip, return_clip)37 38    # ── 2. 涨跌标签 ──────────────────────────────39    df['涨跌标签'] = (df['日收益率'] > 0).astype(int)40 41    # ── 3. 预测目标：下一日涨跌方向（固定1日，与持有期无关） ──42    df['future_ret'] = df['close'].pct_change(periods=forecast_days).shift(-forecast_days)43    df['目标收益率'] = df['future_ret']44    # 训练目标始终是下一日方向45    _next_ret = df['close'].pct_change(periods=1).shift(-1)46    df['目标涨跌'] = _next_ret.gt(0).astype(float)47    df.loc[_next_ret.isna(), '目标涨跌'] = np.nan48 49    # 策略回测用：下一日收益率（T日收盘买入 → T+1日收盘卖出）50    df['next_day_ret'] = df['日收益率'].shift(-1)51 52    # ── 4. 成交量衍生特征 ─────────────────────────53    # 成交量变化率54    df['成交量变化率'] = df['volume'].pct_change()55 56    # 相对成交量（与20日均值比）57    vol_ma20 = df['volume'].rolling(window=20).mean().replace(0, np.nan)58    df['相对成交量'] = df['volume'] / vol_ma2059 60    # 量价配合度（收益率与成交量变化率的滚动相关性）61    df['量价配合度'] = df['日收益率'].rolling(20).corr(df['成交量变化率'])62 63    # 放量上涨: 涨 + 相对成交量 > 1.564    df['放量上涨'] = ((df['日收益率'] > 0) & (df['相对成交量'] > 1.5)).astype(int)65 66    # 缩量下跌: 跌 + 相对成交量 < 0.767    df['缩量下跌'] = ((df['日收益率'] < 0) & (df['相对成交量'] < 0.7)).astype(int)68 69    # ── 5. 多周期动量特征 ────────────────────────────70    df['ret_2d'] = df['close'].pct_change(2)71    df['ret_3d'] = df['close'].pct_change(3)72    df['ret_5d'] = df['close'].pct_change(5)73    df['ret_10d'] = df['close'].pct_change(10)74 75    # ── 6. 均线偏离度 ─────────────────────────────76    ma5 = df['close'].rolling(5).mean()77    ma10 = df['close'].rolling(10).mean()78    ma20 = df['close'].rolling(20).mean()79    df['close_ma5_bias'] = (df['close'] - ma5) / ma580    df['close_ma10_bias'] = (df['close'] - ma10) / ma1081    df['close_ma20_bias'] = (df['close'] - ma20) / ma2082    df['ma5_ma10_cross'] = (ma5 - ma10) / ma1083 84    # ── 7. 波动率特征 ─────────────────────────────85    df['volatility_5d'] = df['日收益率'].rolling(5).std()86    df['volatility_10d'] = df['日收益率'].rolling(10).std()87    df['volatility_20d'] = df['日收益率'].rolling(20).std()88 89    # ── 8. K线形态特征 ─────────────────────────────90    df['upper_shadow'] = (df['high'] - df[['open', 'close']].max(axis=1)) / df['close']91    df['lower_shadow'] = (df[['open', 'close']].min(axis=1) - df['low']) / df['close']92    df['body_ratio'] = (df['close'] - df['open']) / (df['high'] - df['low']).replace(0, np.nan)93    df['high_low_range'] = (df['high'] - df['low']) / df['close']94 95    # ── 9. 多周期 RSI ─────────────────────────────96    for period in [6, 14]:97        delta = df['close'].diff()98        gain = delta.where(delta > 0, 0.0).rolling(period).mean()99        loss = (-delta.where(delta < 0, 0.0)).rolling(period).mean()100        rs = gain / loss.replace(0, np.nan)101        df[f'rsi_{period}'] = 100 - 100 / (1 + rs)102 103    # ── 10. 成交量多周期特征 ──────────────────────────104    df['vol_chg_3d'] = df['volume'].pct_change(3)105    df['vol_chg_5d'] = df['volume'].pct_change(5)106    vol_ma5 = df['volume'].rolling(5).mean()107    df['vol_ratio_5d'] = df['volume'] / vol_ma5.replace(0, np.nan)108 109    # ── 11. 科技股专用特征 ─────────────────────────────110 111    # 换手率系列（科技股情绪核心指标）112    # 仅在 turnover 有效数据超过 50% 时计算，避免 DB 缓存的腾讯源数据全为 NaN113    has_turnover = ('turnover' in df.columns and114                    df['turnover'].notna().mean() > 0.5)115    if has_turnover:116        df['turnover'] = pd.to_numeric(df['turnover'], errors='coerce')117        df['turnover_ma5'] = df['turnover'].rolling(5).mean()118        df['turnover_ma10'] = df['turnover'].rolling(10).mean()119        df['turnover_bias'] = (df['turnover'] - df['turnover_ma5']) / df['turnover_ma5'].replace(0, np.nan)120        df['turnover_accel'] = df['turnover'].diff().rolling(3).mean()121        df['cum_turnover_5d'] = df['turnover'].rolling(5).sum()122        df['cum_turnover_10d'] = df['turnover'].rolling(10).sum()123    elif 'turnover' in df.columns:124        df = df.drop(columns=['turnover'])125 126    # 跳空缺口127    prev_close = df['close'].shift(1)128    df['gap'] = (df['open'] - prev_close) / prev_close129    df['gap_abs'] = df['gap'].abs()130    df['gap_up'] = (df['gap'] > 0.01).astype(int)131    df['gap_down'] = (df['gap'] < -0.01).astype(int)132 133    # 价格位置特征134    for n in [5, 10, 20]:135        rolling_high = df['high'].rolling(n).max()136        rolling_low = df['low'].rolling(n).min()137        price_range = (rolling_high - rolling_low).replace(0, np.nan)138        df[f'price_pos_{n}d'] = (df['close'] - rolling_low) / price_range139        df[f'drawdown_{n}d'] = (df['close'] - rolling_high) / rolling_high140 141    # 连涨连跌天数142    up = (df['日收益率'] > 0).astype(int)143    down = (df['日收益率'] < 0).astype(int)144    streak_up = up.copy()145    streak_down = down.copy()146    for i in range(1, len(df)):147        if up.iloc[i] == 1:148            streak_up.iloc[i] = streak_up.iloc[i - 1] + 1149        if down.iloc[i] == 1:150            streak_down.iloc[i] = streak_down.iloc[i - 1] + 1151    df['streak_up'] = streak_up152    df['streak_down'] = streak_down153 154    # ATR（平均真实波幅）155    tr = pd.concat([156        df['high'] - df['low'],157        (df['high'] - df['close'].shift(1)).abs(),158        (df['low'] - df['close'].shift(1)).abs(),159    ], axis=1).max(axis=1)160    df['atr_5'] = tr.rolling(5).mean()161    df['atr_14'] = tr.rolling(14).mean()162    df['atr_ratio'] = df['atr_5'] / df['atr_14'].replace(0, np.nan)163 164    # 涨停/跌停统计（科技股涨停板效应）165    df['near_limit_up'] = (df['日收益率'] >= 0.09).astype(int)166    df['near_limit_down'] = (df['日收益率'] <= -0.09).astype(int)167    df['limit_up_count_10d'] = df['near_limit_up'].rolling(10).sum()168    df['limit_down_count_10d'] = df['near_limit_down'].rolling(10).sum()169 170    # ── 12. 相对强弱（vs 大盘/板块指数） ─────────────────171    if index_returns is not None:172        idx_ret = index_returns.reindex(df.index)173        df['excess_ret'] = df['日收益率'] - idx_ret.fillna(0)174        df['excess_ret_5d'] = df['excess_ret'].rolling(5).sum()175        df['excess_ret_10d'] = df['excess_ret'].rolling(10).sum()176        df['excess_ret_20d'] = df['excess_ret'].rolling(20).sum()177 178    # ── 13. 处理缺失值 ─────────────────────────────179    # 保留 目标涨跌/future_ret/next_day_ret 为 NaN 的行（最后一天无下日数据），180    # 仅删除特征列（日收益率）为 NaN 的行（首行）181    df = df.dropna(subset=['日收益率'])182 183    return df184 185 186def split_data(df: pd.DataFrame, test_size: float = 0.2) -> dict:187    """188    严格按时间顺序划分训练/测试集（禁止随机）。189 190    返回 dict:191      - train_df, test_df: 完整 DataFrame192      - split_index: 切分点位置193      - latest_close: 最新收盘价（用于反算预测价格）194    """195    split_idx = int(len(df) * (1 - test_size))196    train_df = df.iloc[:split_idx].copy()197    test_df = df.iloc[split_idx:].copy()198 199    return {200        'train_df': train_df,201        'test_df': test_df,202        'split_index': split_idx,203        'latest_close': float(df['close'].iloc[-1]),204    }205 206 207def returns_to_price_series(last_close: float, predicted_returns: np.ndarray,208                            limit_pct: float = None) -> np.ndarray:209    """210    将预测收益率（%）转为预测收盘价序列。211 212    predicted_price[t] = prev_price * (1 + return[t]/100)213    可选应用涨跌幅限制。214    """215    if len(predicted_returns) == 0:216        return np.array([])217 218    prices = np.zeros(len(predicted_returns))219    prev = last_close220    for i, r in enumerate(predicted_returns):221        p = prev * (1 + r / 100)222        if limit_pct is not None:223            upper = last_close * (1 + limit_pct)224            lower = last_close * (1 - limit_pct)225            p = np.clip(p, lower, upper)226        prices[i] = p227        prev = p228    return prices229 230 231def calculate_predicted_prices(last_close: float, predicted_returns: np.ndarray,232                               limit_pct: float = None) -> tuple:233    """234    便捷函数：预测收益率 → 预测收盘价 + 日收益率%。235 236    返回 (predicted_close: np.ndarray, daily_return_pct: np.ndarray)237    """238    prices = returns_to_price_series(last_close, predicted_returns, limit_pct)239    if len(prices) <= 1:240        daily_ret = np.array([])241    else:242        extended = np.concatenate([[last_close], prices])243        daily_ret = (extended[1:] / extended[:-1] - 1) * 100244    return prices, daily_ret