RSI 深度实战：传统信号与 AI 增强结合

RSI = 100 - 100 / (1 + RS)

import yfinance as yf
import numpy as np
import pandas as pd

# ============ 1. 数据准备 ============
def get_data(symbol="AAPL", period="2y"):
    df = yf.download(symbol, period=period, progress=False)
    df = df[["Close"]].copy()
    df.columns = ["close"]
    return df.dropna()

# ============ 2. RSI 计算(纯 pandas,无 talib) ============
def calc_rsi(series: pd.Series, period: int = 14) -> pd.Series:
    delta = series.diff()
    gain = delta.where(delta > 0, 0.0)
    loss = -delta.where(delta < 0, 0.0)
    # Wilder 平滑:等价于 alpha = 1/period 的 EWM
    avg_gain = gain.ewm(alpha=1/period, adjust=False).mean()
    avg_loss = loss.ewm(alpha=1/period, adjust=False).mean()
    rs = avg_gain / avg_loss.replace(0, np.nan)
    rsi = 100 - 100 / (1 + rs)
    return rsi.fillna(50.0)

# ============ 3. 传统 RSI 信号 ============
def rsi_signals(df: pd.DataFrame, lower=30, upper=70) -> pd.DataFrame:
    df = df.copy()
    df["rsi"] = calc_rsi(df["close"])
    df["signal"] = 0
    df.loc[df["rsi"] < lower, "signal"] = 1   # 超卖 -> 买入
    df.loc[df["rsi"] > upper, "signal"] = -1  # 超买 -> 卖出
    return df

# ============ 4. 简单回测:T+1 开盘执行,持有 N 天 ============
def backtest(df: pd.DataFrame, signal_col="signal", hold=5):
    df = df.copy()
    df["ret"] = df["close"].pct_change().fillna(0)
    # 信号次日生效,持有 hold 天后平仓(近似累乘)
    strat_ret = df[signal_col].shift(1) * df["ret"]
    # 累计收益
    cum = (1 + strat_ret.fillna(0)).prod() - 1
    # 胜率:只在有信号的日子统计
    mask = df[signal_col].shift(1).fillna(0) != 0
    wins = (strat_ret[mask] > 0).sum()
    total = mask.sum()
    win_rate = wins / total if total > 0 else 0
    return cum, win_rate, int(total)

if __name__ == "__main__":
    df = get_data("AAPL", "2y")
    df = rsi_signals(df)
    cum, wr, n = backtest(df)
    print(f"[传统 RSI] 累计收益={cum:.2%}, 胜率={wr:.2%}, 信号数={n}")

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report

# ============ 5. 特征工程 ============
def make_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["rsi"] = calc_rsi(df["close"])
    df["rsi_change"] = df["rsi"].diff()
    df["rsi_slope"] = df["rsi"].diff(3)
    df["ret_1d"] = df["close"].pct_change()
    df["ret_3d"] = df["close"].pct_change(3)
    df["dist_ma5"] = df["close"] / df["close"].rolling(5).mean() - 1
    df["dist_ma20"] = df["close"] / df["close"].rolling(20).mean() - 1
    df["volatility"] = df["close"].pct_change().rolling(10).std()
    df["rsi_ob"] = (df["rsi"] > 70).astype(int)
    df["rsi_os"] = (df["rsi"] < 30).astype(int)
    return df

# ============ 6. 样本打标:信号方向 × 未来收益方向 ============
def label_samples(df: pd.DataFrame, forward=5) -> pd.DataFrame:
    df = df.copy()
    future_ret = df["close"].shift(-forward) / df["close"] - 1
    df["label"] = ((df["signal"] * future_ret) > 0).astype(int)
    return df.dropna(subset=["label"])

# ============ 7. 训练 AI 假信号过滤器 ============
def train_filter(df: pd.DataFrame):
    df = label_samples(make_features(df))
    feat_cols = ["rsi", "rsi_change", "rsi_slope", "ret_1d", "ret_3d",
                 "dist_ma5", "dist_ma20", "volatility", "rsi_ob", "rsi_os"]
    # 只在有信号的点取样本
    sig_mask = df["signal"] != 0
    X = df.loc[sig_mask, feat_cols].fillna(0)
    y = df.loc[sig_mask, "label"]
    # 时序切分:前 70% 训练,后 30% 测试(严禁 shuffle)
    split = int(len(X) * 0.7)
    X_tr, X_te = X.iloc[:split], X.iloc[split:]
    y_tr, y_te = y.iloc[:split], y.iloc[split:]
    model = RandomForestClassifier(
        n_estimators=200, max_depth=5, min_samples_leaf=3, random_state=42
    )
    model.fit(X_tr, y_tr)
    pred = model.predict(X_te)
    print(f"AI 过滤器测试集准确率: {accuracy_score(y_te, pred):.2%}")
    print(classification_report(y_te, pred, digits=2))
    return model, feat_cols

# ============ 8. AI 增强回测 ============
def ai_backtest(df: pd.DataFrame, model, feat_cols, hold=5):
    df = make_features(df)
    sig_mask = df["signal"] != 0
    df["ai_signal"] = df["signal"]
    X = df.loc[sig_mask, feat_cols].fillna(0)
    # AI 判断为 1 保留信号,判断为 0 过滤掉
    keep = pd.Series(model.predict(X), index=X.index)
    df.loc[sig_mask, "ai_signal"] = df.loc[sig_mask, "signal"] * keep.values
    return backtest(df, signal_col="ai_signal", hold=hold)

if __name__ == "__main__":
    df = get_data("AAPL", "2y")
    df = rsi_signals(df)
    # 传统 RSI
    cum1, wr1, n1 = backtest(df)
    print(f"[传统 RSI]      收益={cum1:.2%}, 胜率={wr1:.2%}, 次数={n1}")
    # AI 增强
    model, feats = train_filter(df)
    cum2, wr2, n2 = ai_backtest(df, model, feats)
    print(f"[AI 增强 RSI]   收益={cum2:.2%}, 胜率={wr2:.2%}, 次数={n2}")