#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 增强版Kronos预测器 集成量价分析和波浪理论 """ import numpy as np import pandas as pd import torch import torch.nn as nn from typing import Dict, List, Tuple, Optional import warnings warnings.filterwarnings('ignore') # 导入原始Kronos模块 from model.kronos import KronosPredictor, KronosTokenizer, Kronos from model.technical_indicators import TechnicalIndicators, VolumePriceAnalyzer, WaveTheoryAnalyzer, enhance_kronos_features class EnhancedKronosPredictor(KronosPredictor): """增强版Kronos预测器，集成技术指标和波浪理论""" def __init__(self, model, tokenizer, device="cuda:0", max_context=512, clip=5, use_technical_indicators=True, use_wave_theory=True): """ 初始化增强版预测器 Args: model: Kronos模型 tokenizer: Kronos分词器 device: 计算设备 max_context: 最大上下文长度 clip: 数据裁剪值 use_technical_indicators: 是否使用技术指标 use_wave_theory: 是否使用波浪理论 """ super().__init__(model, tokenizer, device, max_context, clip) self.use_technical_indicators = use_technical_indicators self.use_wave_theory = use_wave_theory # 初始化分析器 if self.use_technical_indicators: self.technical_indicators = TechnicalIndicators() self.volume_price_analyzer = VolumePriceAnalyzer() if self.use_wave_theory: self.wave_analyzer = WaveTheoryAnalyzer() # 定义增强特征列表 self.enhanced_features = self._define_enhanced_features() # 更新价格列，包含更多特征 self.price_cols = ['open', 'high', 'low', 'close'] self.vol_col = 'volume' self.amt_vol = 'amount' self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month'] def _define_enhanced_features(self) -> List[str]: """定义增强特征列表""" enhanced_features = [] if self.use_technical_indicators: # 技术指标特征 enhanced_features.extend([ 'sma_5', 'sma_10', 'sma_20', 'sma_50', 'ema_5', 'ema_10', 'ema_20', 'bb_upper', 'bb_middle', 'bb_lower', 'bb_width', 'bb_position', 'k', 'd', 'j', 'rsi_6', 'rsi_14', 'macd', 'macd_signal', 'macd_hist', 'williams_r', 'cci', 'volume_sma_5', 'volume_sma_10', 'volume_sma_20', 'volume_ratio_5', 'volume_ratio_10', 'volume_ratio_20', 'obv', 'ad', 'mfi', 'vwap', 'roc_5', 'roc_10', 'roc_20', 'mom_5', 'mom_10', 'cmo', 'adx', 'plus_di', 'minus_di', 'aroon_up', 'aroon_down', 'aroon_osc', 'sar', 'trix', 'ultosc', 'atr', 'natr', 'trange', 'volatility_5', 'volatility_10', 'volatility_20', 'price_change', 'volume_change', 'price_volume_trend', 'vwap_change', 'volume_price_sync', 'volume_anomaly', 'price_breakout_volume' ]) if self.use_wave_theory: # 波浪理论特征 enhanced_features.extend([ 'wave_count', 'wave_strength', 'fib_5_0.236', 'fib_5_0.382', 'fib_5_0.5', 'fib_5_0.618', 'fib_5_0.786', 'fib_10_0.236', 'fib_10_0.382', 'fib_10_0.5', 'fib_10_0.618', 'fib_10_0.786', 'fib_20_0.236', 'fib_20_0.382', 'fib_20_0.5', 'fib_20_0.618', 'fib_20_0.786', 'fib_50_0.236', 'fib_50_0.382', 'fib_50_0.5', 'fib_50_0.618', 'fib_50_0.786', 'fib_retracement_382', 'fib_retracement_618', 'support_level', 'resistance_level', 'support_distance', 'resistance_distance', 'support_strength', 'resistance_strength', 'wave_momentum_5', 'wave_momentum_10', 'wave_acceleration', 'wave_consistency' ]) return enhanced_features def prepare_enhanced_data(self, df: pd.DataFrame) -> pd.DataFrame: """ 准备增强数据 Args: df: 原始OHLCV数据 Returns: 增强后的数据 """ # 确保数据按时间排序 if 'timestamps' in df.columns: df = df.sort_values('timestamps').reset_index(drop=True) # 计算技术指标和波浪理论特征 enhanced_df = enhance_kronos_features(df) # 处理缺失值 enhanced_df = self._handle_missing_values(enhanced_df) return enhanced_df def _handle_missing_values(self, df: pd.DataFrame) -> pd.DataFrame: """处理缺失值""" # 对于数值列，使用前向填充和后向填充 numeric_columns = df.select_dtypes(include=[np.number]).columns for col in numeric_columns: if df[col].isna().any(): # 前向填充 df[col] = df[col].fillna(method='ffill') # 后向填充 df[col] = df[col].fillna(method='bfill') # 如果还有缺失值，用0填充 df[col] = df[col].fillna(0) return df def predict_enhanced(self, df: pd.DataFrame, x_timestamp: pd.Series, y_timestamp: pd.Series, pred_len: int, T: float = 1.0, top_k: int = 0, top_p: float = 0.9, sample_count: int = 1, verbose: bool = True) -> pd.DataFrame: """ 使用增强特征进行预测 Args: df: 历史数据DataFrame x_timestamp: 历史时间戳 y_timestamp: 预测时间戳 pred_len: 预测长度 T: 采样温度 top_k: top-k采样 top_p: top-p采样 sample_count: 采样次数 verbose: 是否显示详细信息 Returns: 预测结果DataFrame """ if not isinstance(df, pd.DataFrame): raise ValueError("Input must be a pandas DataFrame.") # 检查基础列是否存在 if not all(col in df.columns for col in self.price_cols): raise ValueError(f"Price columns {self.price_cols} not found in DataFrame.") # 准备增强数据 enhanced_df = self.prepare_enhanced_data(df) # 选择用于预测的特征 prediction_features = self._select_prediction_features(enhanced_df) # 处理缺失的volume和amount列 if self.vol_col not in prediction_features.columns: prediction_features[self.vol_col] = 0.0 prediction_features[self.amt_vol] = 0.0 if self.amt_vol not in prediction_features.columns and self.vol_col in prediction_features.columns: prediction_features[self.amt_vol] = prediction_features[self.vol_col] * prediction_features[self.price_cols].mean(axis=1) # 检查是否有NaN值 if prediction_features.isnull().values.any(): print("⚠️ Warning: DataFrame contains NaN values. Filling with zeros.") prediction_features = prediction_features.fillna(0) # 计算时间特征 from model.kronos import calc_time_stamps x_time_df = calc_time_stamps(x_timestamp) y_time_df = calc_time_stamps(y_timestamp) # 准备输入数据 x = prediction_features.values.astype(np.float32) x_stamp = x_time_df.values.astype(np.float32) y_stamp = y_time_df.values.astype(np.float32) # 数据标准化 x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0) x = (x - x_mean) / (x_std + 1e-5) x = np.clip(x, -self.clip, self.clip) # 添加批次维度 x = x[np.newaxis, :] x_stamp = x_stamp[np.newaxis, :] y_stamp = y_stamp[np.newaxis, :] # 生成预测 preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose) # 反标准化 preds = preds.squeeze(0) preds = preds * (x_std + 1e-5) + x_mean # 创建预测结果DataFrame # 只返回OHLCV列 output_columns = self.price_cols + [self.vol_col, self.amt_vol] pred_df = pd.DataFrame(preds[:, :len(output_columns)], columns=output_columns, index=y_timestamp) return pred_df def _select_prediction_features(self, enhanced_df: pd.DataFrame) -> pd.DataFrame: """选择用于预测的特征""" # 基础特征 base_features = self.price_cols + [self.vol_col, self.amt_vol] # 选择存在的增强特征 available_enhanced_features = [f for f in self.enhanced_features if f in enhanced_df.columns] # 组合所有特征 all_features = base_features + available_enhanced_features # 选择存在的特征 existing_features = [f for f in all_features if f in enhanced_df.columns] return enhanced_df[existing_features] def analyze_prediction_quality(self, historical_df: pd.DataFrame, pred_df: pd.DataFrame) -> Dict[str, float]: """ 分析预测质量 Args: historical_df: 历史数据 pred_df: 预测数据 Returns: 预测质量分析结果 """ analysis = {} # 计算历史数据的统计特征 hist_stats = self._calculate_statistics(historical_df) pred_stats = self._calculate_statistics(pred_df) # 计算预测一致性 analysis['price_consistency'] = self._calculate_price_consistency(hist_stats, pred_stats) analysis['volume_consistency'] = self._calculate_volume_consistency(hist_stats, pred_stats) analysis['volatility_consistency'] = self._calculate_volatility_consistency(hist_stats, pred_stats) # 计算趋势一致性 analysis['trend_consistency'] = self._calculate_trend_consistency(historical_df, pred_df) return analysis def _calculate_statistics(self, df: pd.DataFrame) -> Dict[str, float]: """计算数据统计特征""" stats = {} if 'close' in df.columns: stats['close_mean'] = df['close'].mean() stats['close_std'] = df['close'].std() stats['close_volatility'] = df['close'].pct_change().std() if 'volume' in df.columns: stats['volume_mean'] = df['volume'].mean() stats['volume_std'] = df['volume'].std() return stats def _calculate_price_consistency(self, hist_stats: Dict, pred_stats: Dict) -> float: """计算价格一致性""" if 'close_mean' not in hist_stats or 'close_mean' not in pred_stats: return 0.0 # 计算价格变化的合理性 price_change_ratio = abs(pred_stats['close_mean'] - hist_stats['close_mean']) / hist_stats['close_mean'] # 如果价格变化在合理范围内（比如20%以内），认为一致性较高 consistency = max(0, 1 - price_change_ratio / 0.2) return min(1.0, consistency) def _calculate_volume_consistency(self, hist_stats: Dict, pred_stats: Dict) -> float: """计算成交量一致性""" if 'volume_mean' not in hist_stats or 'volume_mean' not in pred_stats: return 0.0 # 计算成交量变化的合理性 volume_change_ratio = abs(pred_stats['volume_mean'] - hist_stats['volume_mean']) / hist_stats['volume_mean'] # 成交量变化在50%以内认为合理 consistency = max(0, 1 - volume_change_ratio / 0.5) return min(1.0, consistency) def _calculate_volatility_consistency(self, hist_stats: Dict, pred_stats: Dict) -> float: """计算波动率一致性""" if 'close_volatility' not in hist_stats or 'close_volatility' not in pred_stats: return 0.0 # 计算波动率变化的合理性 vol_change_ratio = abs(pred_stats['close_volatility'] - hist_stats['close_volatility']) / hist_stats['close_volatility'] # 波动率变化在30%以内认为合理 consistency = max(0, 1 - vol_change_ratio / 0.3) return min(1.0, consistency) def _calculate_trend_consistency(self, historical_df: pd.DataFrame, pred_df: pd.DataFrame) -> float: """计算趋势一致性""" if 'close' not in historical_df.columns or 'close' not in pred_df.columns: return 0.0 # 计算历史趋势 hist_trend = np.polyfit(range(len(historical_df)), historical_df['close'], 1)[0] # 计算预测趋势 pred_trend = np.polyfit(range(len(pred_df)), pred_df['close'], 1)[0] # 计算趋势方向一致性 if hist_trend * pred_trend > 0: # 趋势方向一致 trend_consistency = 1 - abs(hist_trend - pred_trend) / max(abs(hist_trend), abs(pred_trend), 1e-8) else: # 趋势方向相反 trend_consistency = 0 return max(0, trend_consistency) class WaveTheoryEnhancedPredictor(EnhancedKronosPredictor): """专门针对波浪理论优化的预测器""" def __init__(self, model, tokenizer, device="cuda:0", max_context=512, clip=5): super().__init__(model, tokenizer, device, max_context, clip, use_technical_indicators=True, use_wave_theory=True) # 波浪理论特定参数 self.wave_sensitivity = 0.8 # 波浪敏感度 self.fibonacci_weight = 0.3 # 斐波那契权重 self.support_resistance_weight = 0.2 # 支撑阻力权重 def predict_with_wave_analysis(self, df: pd.DataFrame, x_timestamp: pd.Series, y_timestamp: pd.Series, pred_len: int, **kwargs) -> Tuple[pd.DataFrame, Dict]: """ 结合波浪分析进行预测 Args: df: 历史数据 x_timestamp: 历史时间戳 y_timestamp: 预测时间戳 pred_len: 预测长度 **kwargs: 其他预测参数 Returns: (预测结果, 波浪分析结果) """ # 进行波浪分析 wave_analysis = self.wave_analyzer.analyze_wave_patterns(df) # 调整预测参数 based on 波浪分析 adjusted_params = self._adjust_prediction_params(wave_analysis, **kwargs) # 进行预测 pred_df = self.predict_enhanced(df, x_timestamp, y_timestamp, pred_len, **adjusted_params) return pred_df, wave_analysis def _adjust_prediction_params(self, wave_analysis: Dict, **kwargs) -> Dict: """根据波浪分析调整预测参数""" adjusted_params = kwargs.copy() # 根据波浪强度调整采样参数 wave_strength = wave_analysis.get('wave_strength', 0.5) if wave_strength > 0.7: # 强波浪，降低随机性 adjusted_params['T'] = kwargs.get('T', 1.0) * 0.8 adjusted_params['top_p'] = kwargs.get('top_p', 0.9) * 0.95 elif wave_strength < 0.3: # 弱波浪，增加随机性 adjusted_params['T'] = kwargs.get('T', 1.0) * 1.2 adjusted_params['top_p'] = kwargs.get('top_p', 0.9) * 1.05 return adjusted_params def create_enhanced_predictor(model_path: str = "NeoQuasar/Kronos-small", tokenizer_path: str = "NeoQuasar/Kronos-Tokenizer-base", device: str = "cuda:0") -> EnhancedKronosPredictor: """ 创建增强版预测器 Args: model_path: 模型路径 tokenizer_path: 分词器路径 device: 计算设备 Returns: 增强版预测器实例 """ try: # 加载模型和分词器 tokenizer = KronosTokenizer.from_pretrained(tokenizer_path) model = Kronos.from_pretrained(model_path) # 创建增强版预测器 predictor = EnhancedKronosPredictor( model=model, tokenizer=tokenizer, device=device, max_context=512, clip=5, use_technical_indicators=True, use_wave_theory=True ) print("✅ 增强版Kronos预测器创建成功!") print(f" - 技术指标: {'启用' if predictor.use_technical_indicators else '禁用'}") print(f" - 波浪理论: {'启用' if predictor.use_wave_theory else '禁用'}") print(f" - 增强特征数量: {len(predictor.enhanced_features)}") return predictor except Exception as e: print(f"❌ 创建增强版预测器失败: {str(e)}") raise if __name__ == "__main__": # 测试代码 print("测试增强版Kronos预测器...") # 创建示例数据 dates = pd.date_range('2024-01-01', periods=100, freq='D') np.random.seed(42) base_price = 100 prices = [base_price] for i in range(99): change = np.random.normal(0, 0.02) prices.append(max(0.01, prices[-1] * (1 + change))) df = pd.DataFrame({ 'timestamps': dates, 'open': prices, 'high': [p * (1 + abs(np.random.normal(0, 0.01))) for p in prices], 'low': [p * (1 - abs(np.random.normal(0, 0.01))) for p in prices], 'close': prices, 'volume': np.random.normal(1000000, 200000, 100) }) # 测试特征增强 enhanced_df = enhance_kronos_features(df) print(f"原始特征: {len(df.columns)}") print(f"增强特征: {len(enhanced_df.columns)}") print(f"新增特征: {len(enhanced_df.columns) - len(df.columns)}") # 显示部分新特征 new_features = [col for col in enhanced_df.columns if col not in df.columns] print(f"\n新增特征示例:") for feature in new_features[:10]: print(f" {feature}")