Qionk/a-share-quant

1"""2价格预测 - 特征工程3技术指标计算 + 数据归一化 + 滑动窗口4"""5 6import numpy as np7import pandas as pd8from sklearn.preprocessing import MinMaxScaler9 10 11def compute_technical_indicators(df: pd.DataFrame) -> pd.DataFrame:12    """13    计算技术指标，原地添加列14    输入需包含: open, high, low, close, volume15    """16    df = df.copy()17    close = df["close"]18    high = df["high"]19    low = df["low"]20    volume = df["volume"]21 22    # 均线23    for w in [5, 10, 20, 60]:24        df[f"ma{w}"] = close.rolling(w).mean()25 26    # MACD (12, 26, 9)27    ema12 = close.ewm(span=12, adjust=False).mean()28    ema26 = close.ewm(span=26, adjust=False).mean()29    df["dif"] = ema12 - ema2630    df["dea"] = df["dif"].ewm(span=9, adjust=False).mean()31    df["macd"] = (df["dif"] - df["dea"]) * 232 33    # RSI(14)34    delta = close.diff()35    gain = delta.where(delta > 0, 0.0).rolling(14).mean()36    loss = (-delta.where(delta < 0, 0.0)).rolling(14).mean()37    rs = gain / loss.replace(0, np.nan)38    df["rsi"] = 100 - 100 / (1 + rs)39 40    # 布林带 (20, 2)41    ma20 = close.rolling(20).mean()42    std20 = close.rolling(20).std()43    df["boll_upper"] = ma20 + 2 * std2044    df["boll_mid"] = ma2045    df["boll_lower"] = ma20 - 2 * std2046 47    # 涨跌幅（若不存在则计算）48    if "pct_change" not in df.columns or df["pct_change"].isna().all():49        df["pct_change"] = close.pct_change()50 51    # 成交量均线52    df["vol_ma5"] = volume.rolling(5).mean()53    df["vol_ma20"] = volume.rolling(20).mean()54 55    # OBV (On-Balance Volume)56    df["obv"] = (np.sign(close.diff()) * volume).fillna(0).cumsum()57 58    # 量比 (volume / vol_ma20)59    df["vol_ratio"] = volume / df["vol_ma20"].replace(0, np.nan)60 61    # VWAP 近似62    df["vwap"] = (close * volume).cumsum() / volume.cumsum()63 64    # 量价相关性 (20日滚动)65    df["vol_price_corr"] = volume.rolling(20).corr(close)66 67    # 成交量动量 (5日变化率)68    df["vol_momentum"] = volume.pct_change(5)69 70    # 日收益率5日均线（平滑后的收益率趋势）71    if '日收益率' in df.columns:72        df['日收益率_ma5'] = df['日收益率'].rolling(5).mean()73 74    return df75 76 77# 默认特征列（训练用）78DEFAULT_FEATURE_COLS = [79    "目标收益率",  # index 0 = 预测目标80    "close", "open", "high", "low", "volume",81    "ma5", "ma10", "ma20", "ma60",82    "dif", "dea", "macd", "rsi",83    "boll_upper", "boll_mid", "boll_lower",84    "pct_change", "vol_ma5", "vol_ma20",85    "obv", "vol_ratio", "vwap", "vol_price_corr", "vol_momentum",86    "日收益率", "日收益率_ma5",87    "成交量变化率", "相对成交量", "量价配合度",88    "放量上涨", "缩量下跌",89]90 91 92def prepare_features(df: pd.DataFrame, feature_cols: list = None):93    """94    准备特征矩阵：选列 → 去NaN → MinMaxScaler 归一化95    返回: (scaled_array, scaler, feature_cols, cleaned_df)96    """97    if feature_cols is None:98        feature_cols = [c for c in DEFAULT_FEATURE_COLS if c in df.columns]99 100    data = df[feature_cols].copy()101    data = data.replace([np.inf, -np.inf], np.nan)102    data = data.dropna()103 104    scaler = MinMaxScaler()105    scaled = scaler.fit_transform(data.values)106 107    return scaled, scaler, feature_cols, data108 109 110def create_sequences(data: np.ndarray, look_back: int, target_col_idx: int = 0):111    """112    滑动窗口创建训练序列113    data: shape (n_samples, n_features)114    返回: X (n_seq, look_back, n_features), y (n_seq,)115    """116    X, y = [], []117    for i in range(look_back, len(data)):118        X.append(data[i - look_back:i])119        y.append(data[i, target_col_idx])120    return np.array(X), np.array(y)121 122 123def time_series_split(n_samples: int, n_splits: int = 5, min_train_size: int = 100):124    """125    时间序列交叉验证划分（前向扩展）126    返回: [(train_indices, val_indices), ...]127    """128    splits = []129    fold_size = max(1, (n_samples - min_train_size) // (n_splits + 1))130 131    for i in range(n_splits):132        train_end = min_train_size + fold_size * (i + 1)133        # 最后一折的验证集延伸到数据末尾，避免尾部样本不被覆盖134        if i == n_splits - 1:135            val_end = n_samples136        else:137            val_end = min(train_end + fold_size, n_samples)138        if train_end >= n_samples or val_end <= train_end:139            break140        train_idx = np.arange(0, train_end)141        val_idx = np.arange(train_end, val_end)142        splits.append((train_idx, val_idx))143 144    if not splits:145        split_point = int(n_samples * 0.8)146        splits.append((np.arange(0, split_point), np.arange(split_point, n_samples)))147 148    return splits149 150 151def inverse_transform_predictions(predictions: np.ndarray, scaler: MinMaxScaler,152                                   n_features: int, target_col_idx: int = 0) -> np.ndarray:153    """反归一化预测值回原始尺度"""154    dummy = np.zeros((len(predictions), n_features))155    dummy[:, target_col_idx] = predictions.flatten()156    inversed = scaler.inverse_transform(dummy)157    return inversed[:, target_col_idx]158 159 160def create_tabular_features(scaled_data: np.ndarray, feature_cols: list,161                             look_back: int) -> tuple:162    """163    将3D序列数据展平为2D表格特征（用于XGBoost/LightGBM等树模型）。164    每个时间步的特征都作为独立的lag列。165    返回: (X_2D, feature_names)166    """167    n_features = len(feature_cols)168    X, y = [], []169    for i in range(look_back, len(scaled_data)):170        row = scaled_data[i - look_back:i].flatten()171        X.append(row)172    X = np.array(X)173 174    feature_names = []175    for col in feature_cols:176        for lag in range(look_back, 0, -1):177            feature_names.append(f"{col}_lag{lag}")178 179    return X, feature_names180 181 182def create_tabular_targets(scaled_data: np.ndarray, look_back: int,183                            target_col_idx: int = 0) -> np.ndarray:184    """提取表格数据的目标值（归一化后的close价格）"""185    return scaled_data[look_back:, target_col_idx]186