|
|
|
|
| """
|
| 增强版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')
|
|
|
|
|
| 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')
|
|
|
| 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)
|
|
|
|
|
| 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)
|
|
|
|
|
| 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
|
|
|
|
|
|
|
| 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']
|
|
|
|
|
| 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']
|
|
|
|
|
| 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']
|
|
|
|
|
| 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)
|
|
|
|
|
| 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}")
|
|
|
