# ==================== 1. 安装与导入 ====================
!pip install ta catboost lightgbm optuna --quiet
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import warnings
import os
import gc
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import (accuracy_score, roc_auc_score, confusion_matrix,
precision_recall_curve, f1_score)
from sklearn.model_selection import TimeSeriesSplit
import lightgbm as lgb
import catboost as cb
import optuna
import ta
warnings.filterwarnings('ignore')
# ==================== 2. 配置参数 ====================
TICKERS = ['AAPL', 'MSFT', 'NVDA', 'GOOGL', 'AMZN', 'WMT', 'JPM', 'JNJ', 'TMUS', 'AMD']
DATA_PATH = '/kaggle/input/datasets/borismarjanovic/price-volume-data-for-all-us-stocks-etfs/Stocks/'
TRADE_COST = 0.001 # 0.1% 交易成本
results = {}
# ==================== 3. 修复版特征工程（减少NaN）====================
def engineer_features_cls(df):
"""针对分类任务的特征工程，减少NaN产生"""
data = df.copy()
# ----- 基础技术指标（同前）-----
data['SMA_10'] = ta.trend.sma_indicator(data['Close'], 10)
data['SMA_50'] = ta.trend.sma_indicator(data['Close'], 50)
data['EMA_10'] = ta.trend.ema_indicator(data['Close'], 10)
data['EMA_50'] = ta.trend.ema_indicator(data['Close'], 50)
macd = ta.trend.MACD(data['Close'])
data['MACD'] = macd.macd()
data['MACD_signal'] = macd.macd_signal()
data['MACD_diff'] = macd.macd_diff()
data['ADX'] = ta.trend.ADXIndicator(data['High'], data['Low'], data['Close'], 14).adx()
data['RSI'] = ta.momentum.RSIIndicator(data['Close'], 14).rsi()
stoch = ta.momentum.StochasticOscillator(data['High'], data['Low'], data['Close'], 14, 3)
data['Stoch_K'] = stoch.stoch()
data['Stoch_D'] = stoch.stoch_signal()
data['WilliamsR'] = ta.momentum.WilliamsRIndicator(data['High'], data['Low'], data['Close'], 14).williams_r()
bb = ta.volatility.BollingerBands(data['Close'], 20, 2)
data['BB_high'] = bb.bollinger_hband()
data['BB_low'] = bb.bollinger_lband()
data['BB_width'] = bb.bollinger_wband()
data['ATR'] = ta.volatility.AverageTrueRange(data['High'], data['Low'], data['Close'], 14).average_true_range()
data['OBV'] = ta.volume.OnBalanceVolumeIndicator(data['Close'], data['Volume']).on_balance_volume()
data['CMF'] = ta.volume.ChaikinMoneyFlowIndicator(data['High'], data['Low'], data['Close'], data['Volume'], 20).chaikin_money_flow()
data['HL_pct'] = (data['High'] - data['Low']) / data['Close'] * 100
data['OC_pct'] = (data['Close'] - data['Open']) / data['Open'] * 100
data['Volume_Change'] = data['Volume'].pct_change()
# ----- 收益率滞后 -----
for lag in range(1, 6):
data[f'return_lag{lag}'] = data['Close'].pct_change(lag) * 100
# ----- 波动率特征（允许部分NaN） -----
data['vol_5'] = data['Close'].pct_change().rolling(5, min_periods=1).std() * 100
data['vol_20'] = data['Close'].pct_change().rolling(20, min_periods=1).std() * 100
# ----- 窗口比率特征（使用min_periods=1减少NaN）-----
windows = [5, 10, 20]
for w in windows:
# 窗口内均值 / 当前值（使用min_periods=1，使得前期也有值）
data[f'mean_{w}_ratio'] = data['Close'].rolling(w, min_periods=1).mean() / data['Close']
# 窗口内最大值 / 当前值
data[f'max_{w}_ratio'] = data['Close'].rolling(w, min_periods=1).max() / data['Close']
# 窗口内最小值 / 当前值
data[f'min_{w}_ratio'] = data['Close'].rolling(w, min_periods=1).min() / data['Close']
# 成交量均值 / 当前成交量
data[f'volume_mean_{w}_ratio'] = data['Volume'].rolling(w, min_periods=1).mean() / data['Volume']
# ----- 日历特征（不会产生NaN）-----
data['day_of_week'] = data.index.dayofweek
data['month'] = data.index.month
# ----- 交叉特征（可能继承NaN）-----
data['OC_pct_vol'] = data['OC_pct'] * data['Volume_Change']
data['HL_pct_vol'] = data['HL_pct'] * data['Volume_Change']
data['RSI_OC'] = data['RSI'] * data['OC_pct']
# ----- 目标变量（分类）-----
data['target_cls'] = (data['Close'].shift(-1) > data['Close']).astype(int)
return data
# ==================== 4. 定义Transformer模型（用于分类）====================
class TransformerClassifier(nn.Module):
def __init__(self, input_dim, d_model=64, nhead=4, num_layers=2, dim_feedforward=128, dropout=0.1):
super().__init__()
self.input_proj = nn.Linear(input_dim, d_model)
encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.fc = nn.Linear(d_model, 1)
def forward(self, x):
# x shape: (batch, seq_len, input_dim)
x = self.input_proj(x)
x = self.transformer(x)
# 取最后一个时间步的输出
out = self.fc(x[:, -1, :]) # shape: (batch, 1)
return torch.sigmoid(out).squeeze() # shape: (batch,)
# ==================== 5. 特征筛选函数（基于CatBoost重要性）====================
def select_features_cls(X_train, y_train, top_k=200):
"""使用CatBoost分类器筛选特征"""
model = cb.CatBoostClassifier(iterations=100, verbose=0, random_seed=42)
model.fit(X_train, y_train)
importance = model.feature_importances_
feat_imp = pd.DataFrame({'feature': X_train.columns, 'importance': importance})
feat_imp = feat_imp.sort_values('importance', ascending=False).head(top_k)
selected = feat_imp['feature'].tolist()
return selected
# ==================== 6. 单只股票处理函数（分类，无dropna）====================
def process_stock_cls(ticker):
print(f"\n{'='*60}\n处理股票: {ticker}\n{'='*60}")
file_path = os.path.join(DATA_PATH, f"{ticker.lower()}.us.txt")
if not os.path.exists(file_path):
print(f"文件不存在: {file_path}")
return None
df = pd.read_csv(file_path)
df['Date'] = pd.to_datetime(df['Date'])
df = df.set_index('Date').sort_index()
if len(df) < 500:
print(f"原始数据量不足: {len(df)}，跳过")
return None
data = engineer_features_cls(df)
# 处理无穷大（但不dropna）
data = data.replace([np.inf, -np.inf], np.nan)
# 特征列
feature_cols = [c for c in data.columns if c not in ['target_cls']]
X = data[feature_cols]
y = data['target_cls']
# 检查有效样本数（y不能为NaN）
valid_idx = ~y.isna()
X = X[valid_idx]
y = y[valid_idx]
print(f"有效样本数（目标非空）: {len(X)}")
if len(X) < 500:
print(f"有效样本不足: {len(X)}，跳过")
return None
# 划分：训练(60%) + 验证(20%) + 测试(20%)
n = len(X)
train_end = int(n * 0.6)
val_end = int(n * 0.8)
if train_end == 0 or val_end == train_end or n - val_end == 0:
print(f"划分后某集合为空，跳过")
return None
X_train = X.iloc[:train_end]
y_train = y.iloc[:train_end]
X_val = X.iloc[train_end:val_end]
y_val = y.iloc[train_end:val_end]
X_test = X.iloc[val_end:]
y_test = y.iloc[val_end:]
print(f"训练集: {X_train.index[0]} 至 {X_train.index[-1]}, 样本数 {len(X_train)}")
print(f"验证集: {X_val.index[0]} 至 {X_val.index[-1]}, 样本数 {len(X_val)}")
print(f"测试集: {X_test.index[0]} 至 {X_test.index[-1]}, 样本数 {len(X_test)}")
# ---------- 特征筛选 ----------
print("\n正在进行特征筛选...")
if X_train.shape[1] == 0:
print("训练集无特征，跳过")
return None
top_k = min(200, X_train.shape[1])
selected_features = select_features_cls(X_train, y_train, top_k=top_k)
if len(selected_features) == 0:
print("特征筛选后无特征，跳过")
return None
X_train_sel = X_train[selected_features]
X_val_sel = X_val[selected_features]
X_test_sel = X_test[selected_features]
print(f"筛选后特征数: {len(selected_features)}")
# ---------- 模型1: CatBoost ----------
print("\n训练 CatBoost...")
cb_model = cb.CatBoostClassifier(iterations=500, learning_rate=0.05, depth=5,
verbose=0, random_seed=42, eval_metric='AUC')
cb_model.fit(X_train_sel, y_train, eval_set=(X_val_sel, y_val), early_stopping_rounds=50, verbose=False)
cb_pred = cb_model.predict_proba(X_test_sel)[:, 1]
cb_val_pred = cb_model.predict_proba(X_val_sel)[:, 1]
cb_auc = roc_auc_score(y_val, cb_val_pred)
print(f"CatBoost 验证集 AUC: {cb_auc:.4f}")
# ---------- 模型2: LightGBM ----------
print("\n训练 LightGBM...")
lgb_model = lgb.LGBMClassifier(n_estimators=300, max_depth=5, learning_rate=0.05,
random_state=42, verbose=-1)
lgb_model.fit(X_train_sel, y_train, eval_set=[(X_val_sel, y_val)],
callbacks=[lgb.early_stopping(50)])
lgb_pred = lgb_model.predict_proba(X_test_sel)[:, 1]
lgb_val_pred = lgb_model.predict_proba(X_val_sel)[:, 1]
lgb_auc = roc_auc_score(y_val, lgb_val_pred)
print(f"LightGBM 验证集 AUC: {lgb_auc:.4f}")
# ---------- 模型3: Transformer (需构造序列，填充NaN) ----------
seq_len = 20
def create_sequences(X, y, seq_len):
X_seq, y_seq = [], []
for i in range(len(X) - seq_len):
# 获取窗口数据，并将NaN填充为0（简单填充）
window = X.iloc[i:i+seq_len].values
window = np.nan_to_num(window, nan=0.0) # 填充NaN为0
X_seq.append(window)
y_seq.append(y.iloc[i+seq_len])
return np.array(X_seq), np.array(y_seq)
# 检查数据量是否足够构造序列
if len(X_train_sel) <= seq_len or len(X_val_sel) <= seq_len or len(X_test_sel) <= seq_len:
print("数据量不足以构造序列，跳过Transformer")
# 只用树模型融合
weights = np.array([cb_auc, lgb_auc]) / (cb_auc + lgb_auc)
ensemble_val_prob = weights[0]*cb_val_pred + weights[1]*lgb_val_pred
ensemble_val_auc = roc_auc_score(y_val, ensemble_val_prob)
print(f"仅树模型融合验证集 AUC: {ensemble_val_auc:.4f}")
ensemble_test_prob = weights[0]*cb_pred + weights[1]*lgb_pred
ensemble_test_auc = roc_auc_score(y_test, ensemble_test_prob)
print(f"仅树模型融合测试集 AUC: {ensemble_test_auc:.4f}")
# 阈值优化
precision, recall, thresholds = precision_recall_curve(y_val, ensemble_val_prob)
f1_scores = 2 * (precision * recall) / (precision + recall + 1e-10)
best_idx = np.argmax(f1_scores[:-1])
best_thresh = thresholds[best_idx]
test_pred_label = (ensemble_test_prob >= best_thresh).astype(int)
test_acc = accuracy_score(y_test, test_pred_label)
# 策略回测
test_df = X_test_sel.copy()
test_df['actual_return'] = df.loc[test_df.index, 'Close'].pct_change().shift(-1).values * 100
test_df['pred_prob'] = ensemble_test_prob
test_df['signal'] = test_pred_label
test_df['prev_signal'] = test_df['signal'].shift(1).fillna(0)
test_df['trade'] = (test_df['signal'] != test_df['prev_signal']).astype(int)
test_df['strategy_return'] = test_df['signal'] * test_df['actual_return'] - test_df['trade'] * TRADE_COST * 100
test_df['buy_hold_return'] = test_df['actual_return']
test_df['strategy_cum'] = (1 + test_df['strategy_return']/100).cumprod()
test_df['buy_hold_cum'] = (1 + test_df['buy_hold_return']/100).cumprod()
strategy_return = (test_df['strategy_cum'].iloc[-1] - 1) * 100
buy_hold_return = (test_df['buy_hold_cum'].iloc[-1] - 1) * 100
sharpe = test_df['strategy_return'].mean() / test_df['strategy_return'].std() * np.sqrt(252)
return {
'ticker': ticker,
'test_auc': ensemble_test_auc,
'test_acc': test_acc,
'best_thresh': best_thresh,
'strategy_return': strategy_return,
'buy_hold_return': buy_hold_return,
'sharpe': sharpe,
'n_trades': test_df['trade'].sum()
}
X_train_seq, y_train_seq = create_sequences(X_train_sel, y_train, seq_len)
X_val_seq, y_val_seq = create_sequences(X_val_sel, y_val, seq_len)
X_test_seq, y_test_seq = create_sequences(X_test_sel, y_test, seq_len)
print(f"\nTransformer 训练样本数: {len(X_train_seq)}")
# 转换为PyTorch张量
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_dataset = TensorDataset(torch.FloatTensor(X_train_seq), torch.FloatTensor(y_train_seq))
val_dataset = TensorDataset(torch.FloatTensor(X_val_seq), torch.FloatTensor(y_val_seq))
test_dataset = TensorDataset(torch.FloatTensor(X_test_seq), torch.FloatTensor(y_test_seq))
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
# 初始化Transformer
input_dim = X_train_sel.shape[1]
tf_model = TransformerClassifier(input_dim=input_dim, d_model=64, nhead=4, num_layers=2).to(device)
criterion = nn.BCELoss()
optimizer = optim.Adam(tf_model.parameters(), lr=0.001)
# 训练（早停）
best_val_auc = 0
patience = 10
counter = 0
for epoch in range(100):
tf_model.train()
train_loss = 0
for Xb, yb in train_loader:
Xb, yb = Xb.to(device), yb.to(device)
optimizer.zero_grad()
pred = tf_model(Xb)
loss = criterion(pred, yb)
loss.backward()
optimizer.step()
train_loss += loss.item()
# 验证
tf_model.eval()
val_preds = []
with torch.no_grad():
for Xb, _ in val_loader:
Xb = Xb.to(device)
pred = tf_model(Xb)
val_preds.extend(pred.cpu().numpy().flatten().tolist()) # 修复：确保一维
val_auc = roc_auc_score(y_val_seq, val_preds)
if val_auc > best_val_auc:
best_val_auc = val_auc
counter = 0
torch.save(tf_model.state_dict(), 'best_tf.pth')
else:
counter += 1
if counter >= patience:
break
# 加载最佳模型并在测试集上预测
tf_model.load_state_dict(torch.load('best_tf.pth'))
tf_model.eval()
tf_test_preds = []
with torch.no_grad():
for Xb, _ in test_loader:
Xb = Xb.to(device)
pred = tf_model(Xb)
tf_test_preds.extend(pred.cpu().numpy().flatten().tolist()) # 修复
tf_test_preds = np.array(tf_test_preds)
# 对齐标签
y_test_aligned = y_test.iloc[seq_len:].values
y_val_aligned = y_val.iloc[seq_len:].values
# 对齐树模型预测
cb_val_aligned = cb_val_pred[seq_len:]
lgb_val_aligned = lgb_val_pred[seq_len:]
cb_test_aligned = cb_pred[seq_len:]
lgb_test_aligned = lgb_pred[seq_len:]
# ---------- 模型融合 ----------
aucs = np.array([roc_auc_score(y_val_aligned, cb_val_aligned),
roc_auc_score(y_val_aligned, lgb_val_aligned),
best_val_auc])
weights = aucs / aucs.sum()
print(f"\n融合权重: CB={weights[0]:.3f}, LGB={weights[1]:.3f}, TF={weights[2]:.3f}")
ensemble_val_prob = weights[0]*cb_val_aligned + weights[1]*lgb_val_aligned + weights[2]*np.array(val_preds)
ensemble_val_auc = roc_auc_score(y_val_aligned, ensemble_val_prob)
print(f"融合模型验证集 AUC: {ensemble_val_auc:.4f}")
ensemble_test_prob = weights[0]*cb_test_aligned + weights[1]*lgb_test_aligned + weights[2]*tf_test_preds
ensemble_test_auc = roc_auc_score(y_test_aligned, ensemble_test_prob)
print(f"融合模型测试集 AUC: {ensemble_test_auc:.4f}")
# ---------- 阈值优化与策略回测 ----------
precision, recall, thresholds = precision_recall_curve(y_val_aligned, ensemble_val_prob)
f1_scores = 2 * (precision * recall) / (precision + recall + 1e-10)
best_idx = np.argmax(f1_scores[:-1])
best_thresh = thresholds[best_idx]
print(f"最佳阈值 (基于验证集F1): {best_thresh:.4f}")
test_pred_label = (ensemble_test_prob >= best_thresh).astype(int)
test_acc = accuracy_score(y_test_aligned, test_pred_label)
print(f"测试集准确率: {test_acc:.4f}")
test_df = X_test_sel.iloc[seq_len:].copy()
test_df['actual_return'] = df.loc[test_df.index, 'Close'].pct_change().shift(-1).values * 100
test_df['pred_prob'] = ensemble_test_prob
test_df['signal'] = test_pred_label
test_df['prev_signal'] = test_df['signal'].shift(1).fillna(0)
test_df['trade'] = (test_df['signal'] != test_df['prev_signal']).astype(int)
test_df['strategy_return'] = test_df['signal'] * test_df['actual_return'] - test_df['trade'] * TRADE_COST * 100
test_df['buy_hold_return'] = test_df['actual_return']
test_df['strategy_cum'] = (1 + test_df['strategy_return']/100).cumprod()
test_df['buy_hold_cum'] = (1 + test_df['buy_hold_return']/100).cumprod()
strategy_return = (test_df['strategy_cum'].iloc[-1] - 1) * 100
buy_hold_return = (test_df['buy_hold_cum'].iloc[-1] - 1) * 100
sharpe = test_df['strategy_return'].mean() / test_df['strategy_return'].std() * np.sqrt(252)
print(f"\n策略总收益率 (考虑成本): {strategy_return:.2f}%")
print(f"买入持有收益率: {buy_hold_return:.2f}%")
print(f"策略夏普比率: {sharpe:.2f}")
print(f"交易次数: {test_df['trade'].sum()}")
plt.figure(figsize=(10,5))
plt.plot(test_df.index, test_df['strategy_cum'], label='Strategy (with cost)')
plt.plot(test_df.index, test_df['buy_hold_cum'], label='Buy & Hold')
plt.title(f'{ticker} - Classification Strategy vs Buy & Hold')
plt.xlabel('Date')
plt.ylabel('Cumulative Return')
plt.legend()
plt.grid(True)
plt.show()
return {
'ticker': ticker,
'test_auc': ensemble_test_auc,
'test_acc': test_acc,
'best_thresh': best_thresh,
'strategy_return': strategy_return,
'buy_hold_return': buy_hold_return,
'sharpe': sharpe,
'n_trades': test_df['trade'].sum()
}
# ==================== 7. 主循环 ====================
for ticker in TICKERS:
result = process_stock_cls(ticker)
if result:
results[ticker] = result
gc.collect()
# ==================== 8. 结果汇总 ====================
valid_results = {k: v for k, v in results.items() if v is not None and isinstance(v, dict)}
if len(valid_results) == 0:
print("\n没有有效结果，请检查各股票处理是否成功。")
else:
summary = pd.DataFrame(valid_results).T
print("\n所有结果包含的键：", set().union(*[set(v.keys()) for v in valid_results.values()]))
required_cols = ['test_auc', 'test_acc', 'strategy_return', 'buy_hold_return', 'sharpe', 'n_trades']
missing_cols = [col for col in required_cols if col not in summary.columns]
if missing_cols:
print(f"\n警告：下列列缺失：{missing_cols}，无法生成汇总。")
else:
print("\n\n========== 多股票分类结果汇总 ==========")
print(summary[required_cols].to_string())
fig, axes = plt.subplots(2, 2, figsize=(14,10))
summary['test_auc'].sort_values().plot(kind='bar', ax=axes[0,0], color='skyblue')
axes[0,0].set_title('Test AUC')
axes[0,0].axhline(y=0.5, color='red', linestyle='--')
summary['test_acc'].sort_values().plot(kind='bar', ax=axes[0,1], color='lightgreen')
axes[0,1].set_title('Test Accuracy')
axes[0,1].axhline(y=0.5, color='red', linestyle='--')
summary[['strategy_return', 'buy_hold_return']].plot(kind='bar', ax=axes[1,0])
axes[1,0].set_title('Strategy vs Buy & Hold Return')
summary['sharpe'].sort_values().plot(kind='bar', ax=axes[1,1], color='coral')
axes[1,1].set_title('Strategy Sharpe Ratio')
plt.tight_layout()
plt.show()
summary.to_csv('classification_results.csv')
print("\n结果已保存至 classification_results.csv")
