Qionk/a-share-quant

1"""2价格预测 - 模型定义3LSTM / GRU / 1D-CNN / CNN-GRU / PatchTST / TFT /4XGBoost / LightGBM / ARIMA / SARIMA / GARCH5"""6 7import os8import numpy as np9 10os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"11 12from dataclasses import dataclass, field13 14 15@dataclass16class ModelConfig:17    """模型超参数"""18    look_back: int = 3019    n_features: int = 120    epochs: int = 10021    batch_size: int = 3222    early_stop_patience: int = 1023    learning_rate: float = 0.00124    dropout: float = 0.225    lstm_units: list = field(default_factory=lambda: [64, 32])26    gru_units: list = field(default_factory=lambda: [64, 32])27    cnn_filters: list = field(default_factory=lambda: [64, 32])28    cnn_kernel_size: int = 329    # PatchTST30    patchtst_patch_size: int = 1631    patchtst_d_model: int = 12832    patchtst_n_heads: int = 433    patchtst_n_encoder_layers: int = 234    patchtst_ff_dim: int = 25635    patchtst_dropout: float = 0.136    # TFT37    tft_hidden_size: int = 6438    tft_n_heads: int = 439    tft_dropout: float = 0.240    tft_lstm_layers: int = 141    # CNN-GRU42    cnn_gru_filters: list = field(default_factory=lambda: [32, 16])43    cnn_gru_gru_units: list = field(default_factory=lambda: [24])44    cnn_gru_kernel_size: int = 345    # XGBoost46    xgboost_n_estimators: int = 10047    xgboost_max_depth: int = 648    xgboost_learning_rate: float = 0.149    xgboost_subsample: float = 0.850    # LightGBM51    lightgbm_n_estimators: int = 10052    lightgbm_max_depth: int = 653    lightgbm_learning_rate: float = 0.154    lightgbm_num_leaves: int = 3155    lightgbm_subsample: float = 0.856    # Per-model DL learning rates57    lstm_lr: float = 0.00158    gru_lr: float = 0.00159    cnn_lr: float = 0.00160    cnn_gru_lr: float = 0.00161    patchtst_lr: float = 0.00162    tft_lr: float = 0.00163    # SARIMA64    sarima_order: tuple = (1, 1, 1)65    sarima_seasonal_order: tuple = (1, 1, 1, 5)66    # GARCH67    garch_p: int = 168    garch_q: int = 169    garch_dist: str = 't'70 71 72def _set_seed(seed=42):73    import tensorflow as tf74    np.random.seed(seed)75    tf.random.set_seed(seed)76 77 78def build_lstm(config: ModelConfig):79    """80    LSTM 模型:81    LSTM(64, return_sequences) → Dropout → LSTM(32) → Dropout → Dense(1)82    """83    _set_seed()84    import tensorflow as tf85    from tensorflow.keras.models import Sequential86    from tensorflow.keras.layers import LSTM, Dense, Dropout, Input87 88    model = Sequential([89        Input(shape=(config.look_back, config.n_features)),90        LSTM(config.lstm_units[0], return_sequences=True),91        Dropout(config.dropout),92        LSTM(config.lstm_units[1]),93        Dropout(config.dropout),94        Dense(16, activation="relu"),95        Dense(1),96    ])97    model.compile(98        optimizer=tf.keras.optimizers.Adam(learning_rate=config.lstm_lr),99        loss="mse",100    )101    return model102 103 104def build_gru(config: ModelConfig):105    """106    GRU 模型: 与 LSTM 同构，替换为 GRU 层107    """108    _set_seed()109    import tensorflow as tf110    from tensorflow.keras.models import Sequential111    from tensorflow.keras.layers import GRU, Dense, Dropout, Input112 113    model = Sequential([114        Input(shape=(config.look_back, config.n_features)),115        GRU(config.gru_units[0], return_sequences=True),116        Dropout(config.dropout),117        GRU(config.gru_units[1]),118        Dropout(config.dropout),119        Dense(16, activation="relu"),120        Dense(1),121    ])122    model.compile(123        optimizer=tf.keras.optimizers.Adam(learning_rate=config.gru_lr),124        loss="mse",125    )126    return model127 128 129def build_cnn(config: ModelConfig):130    """131    1D-CNN 模型:132    Conv1D(64,3) → MaxPool → Conv1D(32,3) → MaxPool → Flatten → Dense(32) → Dense(1)133    """134    _set_seed()135    import tensorflow as tf136    from tensorflow.keras.models import Sequential137    from tensorflow.keras.layers import (138        Conv1D, MaxPooling1D, Flatten, Dense, Dropout, Input139    )140 141    layers = [Input(shape=(config.look_back, config.n_features))]142 143    seq_len = config.look_back144    layers.append(Conv1D(config.cnn_filters[0], config.cnn_kernel_size,145                         activation="relu", padding="same"))146    if seq_len >= 4:147        layers.append(MaxPooling1D(pool_size=2))148        seq_len = seq_len // 2149 150    layers.append(Conv1D(config.cnn_filters[1], config.cnn_kernel_size,151                         activation="relu", padding="same"))152    if seq_len >= 4:153        layers.append(MaxPooling1D(pool_size=2))154 155    layers.extend([156        Flatten(),157        Dense(32, activation="relu"),158        Dropout(config.dropout),159        Dense(1),160    ])161 162    model = Sequential(layers)163    model.compile(164        optimizer=tf.keras.optimizers.Adam(learning_rate=config.cnn_lr),165        loss="mse",166    )167    return model168 169 170def build_cnn_gru(config: ModelConfig):171    """172    CNN-GRU 混合模型（优化版）:173    Conv1D(32) → Conv1D(16) → GRU(24) → Dropout → Dense(1)174    移除 MaxPooling 防信息丢失，单层 GRU 减过拟合，去除冗余 Dense 层175    """176    _set_seed()177    import tensorflow as tf178    from tensorflow.keras.models import Sequential179    from tensorflow.keras.layers import (180        Conv1D, GRU, Dense, Dropout, Input181    )182 183    layers = [Input(shape=(config.look_back, config.n_features))]184    layers.append(Conv1D(config.cnn_gru_filters[0], config.cnn_gru_kernel_size,185                         activation="relu", padding="valid"))186    layers.append(Conv1D(config.cnn_gru_filters[1], config.cnn_gru_kernel_size,187                         activation="relu", padding="valid"))188    layers.append(GRU(config.cnn_gru_gru_units[0], return_sequences=False))189    layers.append(Dropout(config.dropout))190    layers.append(Dense(1))191 192    model = Sequential(layers)193    model.compile(194        optimizer=tf.keras.optimizers.Adam(learning_rate=config.cnn_gru_lr),195        loss="mse",196    )197    return model198 199 200def build_patchtst(config: ModelConfig):201    """202    PatchTST + RevIN: per-instance z-score → Patch → Transformer → inverse203    """204    _set_seed()205    import tensorflow as tf206    from tensorflow.keras import layers, Model207 208    look_back = config.look_back209    n_features = config.n_features210    patch_size = config.patchtst_patch_size211    d_model = config.patchtst_d_model212    n_heads = config.patchtst_n_heads213    n_layers = config.patchtst_n_encoder_layers214    ff_dim = config.patchtst_ff_dim215    drop = config.patchtst_dropout216 217    while look_back // patch_size < 3 and patch_size > 4:218        patch_size = patch_size // 2219 220    n_patches = (look_back + patch_size - 1) // patch_size221    pad_len = n_patches * patch_size - look_back222 223    inputs = layers.Input(shape=(look_back, n_features))224 225    # ── RevIN normalize: (x - mean) / std ──226    eps = 1e-5227    x_mean = layers.Lambda(lambda t: tf.reduce_mean(t, axis=1, keepdims=True))(inputs)228    x_std = layers.Lambda(lambda t: tf.math.reduce_std(t, axis=1, keepdims=True) + 1e-5)(inputs)229    x_norm = layers.Subtract()([inputs, x_mean])230    x_norm = layers.Lambda(lambda t: t[0] / t[1])([x_norm, x_std])231 232    # Save target feature stats for inverse (reshape to (None,1) for broadcasting)233    target_mean = layers.Lambda(lambda t: t[:, 0, 0:1])(x_mean)234    target_std = layers.Lambda(lambda t: t[:, 0, 0:1])(x_std)235 236    x = x_norm237    if pad_len > 0:238        x = layers.ZeroPadding1D(padding=(pad_len, 0))(x)239 240    x = layers.Reshape((n_patches, patch_size * n_features))(x)241    x = layers.Dense(d_model)(x)242 243    pos_emb = layers.Embedding(n_patches, d_model)244    positions = tf.range(n_patches)245    x = x + pos_emb(positions)246 247    for _ in range(n_layers):248        attn_out = layers.MultiHeadAttention(249            num_heads=n_heads, key_dim=d_model // n_heads, dropout=drop250        )(x, x)251        attn_out = layers.Dropout(drop)(attn_out)252        x = layers.LayerNormalization(epsilon=1e-6)(x + attn_out)253 254        ff_out = layers.Dense(ff_dim, activation="gelu")(x)255        ff_out = layers.Dropout(drop)(ff_out)256        ff_out = layers.Dense(d_model)(ff_out)257        ff_out = layers.Dropout(drop)(ff_out)258        x = layers.LayerNormalization(epsilon=1e-6)(x + ff_out)259 260    x = layers.Flatten()(x)261    x = layers.Dense(64, activation="relu")(x)262    x = layers.Dropout(drop)(x)263    x = layers.Dense(1)(x)264 265    # ── RevIN inverse: x * std + mean ──266    outputs = layers.Multiply()([x, target_std])267    outputs = layers.Add()([outputs, target_mean])268 269    model = Model(inputs=inputs, outputs=outputs, name="PatchTST")270    model.compile(271        optimizer=tf.keras.optimizers.Adam(learning_rate=config.patchtst_lr),272        loss="mse",273    )274    return model275 276 277def build_tft(config: ModelConfig):278    """279    Temporal Fusion Transformer (简化版):280    Variable Selection → LSTM Encoder → GRN → Temporal Attention → Output281    """282    _set_seed()283    import tensorflow as tf284    from tensorflow.keras import layers, Model285 286    look_back = config.look_back287    n_features = config.n_features288    hidden = config.tft_hidden_size289    n_heads = config.tft_n_heads290    drop = config.tft_dropout291    n_lstm = config.tft_lstm_layers292 293    inputs = layers.Input(shape=(look_back, n_features))294 295    # Variable Selection Network296    context = layers.Flatten()(inputs)297    context = layers.Dense(hidden, activation="relu")(context)298    var_weights = layers.Dense(n_features, activation="softmax",299                               name="variable_weights")(context)300 301    # Per-feature projection + weighted sum302    split_proj = []303    for i in range(n_features):304        feat_slice = layers.Lambda(lambda x, idx=i: x[:, :, idx:idx+1])(inputs)305        feat_proj = layers.Dense(hidden)(feat_slice)306        split_proj.append(feat_proj)307 308    stacked = layers.Lambda(lambda x: tf.stack(x, axis=2))(split_proj)309 310    var_w_expanded = layers.Lambda(311        lambda x: tf.expand_dims(tf.expand_dims(x, 1), -1)312    )(var_weights)313 314    weighted = layers.Multiply()([stacked, var_w_expanded])315    selected = layers.Lambda(lambda x: tf.reduce_sum(x, axis=2))(weighted)316 317    # LSTM Encoder318    lstm_out = selected319    for i in range(n_lstm):320        lstm_out = layers.LSTM(hidden, return_sequences=True,321                               dropout=drop, name=f"tft_lstm_{i}")(lstm_out)322 323    # Gated Residual Network324    grn_h = layers.Dense(hidden, activation="elu")(lstm_out)325    grn_h = layers.Dense(hidden)(grn_h)326    grn_h = layers.Dropout(drop)(grn_h)327    gate = layers.Dense(hidden, activation="sigmoid")(lstm_out)328    grn_out = layers.Multiply()([gate, grn_h])329    skip = layers.Lambda(lambda x: (1 - x[0]) * x[1])([gate, lstm_out])330    grn_out = layers.Add()([grn_out, skip])331    grn_out = layers.LayerNormalization()(grn_out)332 333    # Temporal Self-Attention with static enrichment334    static_ctx = layers.GlobalAveragePooling1D()(lstm_out)335    static_ctx = layers.RepeatVector(look_back)(static_ctx)336    enriched = layers.Add()([grn_out, static_ctx])337 338    attn_out = layers.MultiHeadAttention(339        num_heads=n_heads, key_dim=hidden // n_heads, dropout=drop340    )(enriched, enriched)341    attn_out = layers.Dropout(drop)(attn_out)342    attn_out = layers.Add()([attn_out, enriched])343    attn_out = layers.LayerNormalization()(attn_out)344 345    # Output346    last_step = layers.Lambda(lambda x: x[:, -1, :])(attn_out)347    out = layers.Dense(hidden // 2, activation="relu")(last_step)348    out = layers.Dropout(drop)(out)349    outputs = layers.Dense(1)(out)350 351    model = Model(inputs=inputs, outputs=outputs, name="TFT")352    model.compile(353        optimizer=tf.keras.optimizers.Adam(learning_rate=config.tft_lr),354        loss="mse",355    )356    return model357 358 359def fit_arima(train_data: np.ndarray, exog: np.ndarray = None, config: ModelConfig = None):360    """361    拟合 ARIMA(1,1,1) + 常数漂移362    train_data: 一维收盘价序列363    """364    from statsmodels.tsa.arima.model import ARIMA365    import warnings366    warnings.filterwarnings("ignore")367 368    model = ARIMA(train_data, order=(1, 1, 1), exog=exog, trend='t')369    result = model.fit()370    return result371 372 373def predict_arima(model, steps: int, exog_future: np.ndarray = None):374    """375    ARIMA 多步预测376    返回: (predictions, confidence_intervals)377    """378    forecast = model.get_forecast(steps=steps, exog=exog_future)379    fc = np.array(forecast.predicted_mean)380    conf = np.array(forecast.conf_int(alpha=0.05))381    return fc, conf382 383 384def returns_to_prices(last_price: float, predicted_returns: np.ndarray) -> np.ndarray:385    """从最后收盘价 + 预测收益率序列 → 价格序列"""386    prices = [last_price]387    for r in predicted_returns:388        prices.append(prices[-1] * (1 + r / 100))389    return np.array(prices[1:])390 391 392def build_xgboost(config: ModelConfig):393    """XGBoost 回归模型"""394    import xgboost as xgb395    return xgb.XGBRegressor(396        n_estimators=config.xgboost_n_estimators,397        max_depth=config.xgboost_max_depth,398        learning_rate=config.xgboost_learning_rate,399        subsample=config.xgboost_subsample,400        random_state=42,401        n_jobs=-1,402        verbosity=0,403    )404 405 406def build_lightgbm(config: ModelConfig):407    """LightGBM 回归模型"""408    import lightgbm as lgb409    return lgb.LGBMRegressor(410        n_estimators=config.lightgbm_n_estimators,411        max_depth=config.lightgbm_max_depth,412        learning_rate=config.lightgbm_learning_rate,413        num_leaves=config.lightgbm_num_leaves,414        random_state=42,415        n_jobs=-1,416        verbose=-1,417    )418 419 420def fit_sarima(train_data: np.ndarray, config: ModelConfig = None):421    """422    拟合 SARIMA 模型423    train_data: 一维收盘价序列424    """425    from statsmodels.tsa.statespace.sarimax import SARIMAX426    import warnings427    warnings.filterwarnings("ignore")428 429    order = config.sarima_order if config else (1, 1, 1)430    seasonal_order = config.sarima_seasonal_order if config else (1, 1, 1, 5)431 432    model = SARIMAX(train_data, order=order, seasonal_order=seasonal_order,433                    trend='t', enforce_stationarity=False, enforce_invertibility=False)434    result = model.fit(disp=False)435    return result436 437 438def predict_sarima(model, steps: int):439    """440    SARIMA 多步预测441    返回: (predictions, confidence_intervals)442    """443    forecast = model.get_forecast(steps=steps)444    fc = np.array(forecast.predicted_mean)445    conf = np.array(forecast.conf_int(alpha=0.05))446    return fc, conf447