import gc import os import re import subprocess import time from datetime import datetime, timedelta from pathlib import Path import matplotlib.pyplot as plt import numpy as np import pandas as pd import torch import baostock as bs from pytz import timezone # 处理中国时区（Asia/Shanghai） from model import KronosTokenizer, Kronos, KronosPredictor # --- Configuration --- Config = { "REPO_PATH": Path(__file__).parent.resolve(), "LOCAL_MODEL_PATH": os.path.join(Path(__file__).parent.resolve(), "models"), "STOCK_CODE": "sh.000001", "FREQUENCY": "d", "START_DATE": "2022-01-01", "PRED_HORIZON": 24, "N_PREDICTIONS": 10, "VOL_WINDOW": 24, "PREDICTION_CACHE": os.path.join("/tmp", "predictions_cache"), "CHART_PATH": os.path.join("/tmp", "prediction_chart.png"), "HTML_PATH": os.path.join("/tmp", "index.html"), # 先不定义CHINESE_FONT_PATH，避免引用未完成的Config "IS_TODAY_INFERENCED": False, "CACHED_RESULTS": { "close_preds": None, "volume_preds": None, "v_close_preds": None, "upside_prob": None, "vol_amp_prob": None, "hist_df_for_plot": None } } # 补充定义中文字体路径（此时Config已完全定义） Config["CHINESE_FONT_PATH"] = os.path.join(Config["REPO_PATH"], "fonts", "wqy-microhei.ttf") # 创建必要目录 os.makedirs(Config["PREDICTION_CACHE"], exist_ok=True) os.makedirs(Config["LOCAL_MODEL_PATH"], exist_ok=True) def get_china_time(): """获取当前中国时间（Asia/Shanghai时区），返回datetime对象""" china_tz = timezone("Asia/Shanghai") return datetime.now(china_tz) def load_local_model(): """加载本地Kronos模型，添加字体加载日志""" print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 开始加载本地Kronos模型...") tokenizer_path = os.path.join(Config["LOCAL_MODEL_PATH"], "tokenizer") model_path = os.path.join(Config["LOCAL_MODEL_PATH"], "model") # 检查模型文件是否存在 if not os.path.exists(tokenizer_path): raise FileNotFoundError(f"分词器路径不存在：{tokenizer_path}") if not os.path.exists(model_path): raise FileNotFoundError(f"模型路径不存在：{model_path}") # 加载模型和分词器 tokenizer = KronosTokenizer.from_pretrained(tokenizer_path, local_files_only=True) model = Kronos.from_pretrained(model_path, local_files_only=True) tokenizer.eval() model.eval() predictor = KronosPredictor(model, tokenizer, device="cpu", max_context=512) print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 本地模型加载成功") return predictor def fetch_stock_data(): """获取股票数据（每日更新一次，中国时间），添加数据获取日志""" china_now = get_china_time() end_date = china_now.strftime("%Y-%m-%d") # 按中国时间取结束日期 need_points = Config["VOL_WINDOW"] + Config["VOL_WINDOW"] # 历史数据+波动率计算窗口 print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 开始获取{Config['STOCK_CODE']}日线数据（结束日期：{end_date}）") lg = bs.login() if lg.error_code != '0': raise ConnectionError(f"Baostock登录失败：{lg.error_msg}") try: # 调用baostock获取K线数据 fields = "date,open,high,low,close,volume" rs = bs.query_history_k_data_plus( code=Config["STOCK_CODE"], fields=fields, start_date=Config["START_DATE"], end_date=end_date, frequency=Config["FREQUENCY"], adjustflag="2" # 后复权 ) if rs.error_code != '0': raise ValueError(f"获取K线数据失败：{rs.error_msg}") # 处理数据 data_list = [] while rs.next(): data_list.append(rs.get_row_data()) df = pd.DataFrame(data_list, columns=rs.fields) # 数值列转换 numeric_cols = ['open', 'high', 'low', 'close', 'volume'] for col in numeric_cols: df[col] = pd.to_numeric(df[col], errors='coerce') df = df.dropna(subset=numeric_cols) # 添加时间戳和成交额列 df['timestamps'] = pd.to_datetime(df['date'], format='%Y-%m-%d') df['amount'] = (df['open'] + df['high'] + df['low'] + df['close']) / 4 * df['volume'] df = df[['timestamps', 'open', 'high', 'low', 'close', 'volume', 'amount']] # 检查数据量 if len(df) < need_points: raise ValueError(f"数据量不足（仅{len(df)}个交易日），请提前START_DATE") df = df.tail(need_points).reset_index(drop=True) print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 股票数据获取成功，共{len(df)}个交易日") print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 最新5条数据：\n{df[['timestamps', 'open', 'close', 'volume']].tail()}") return df finally: bs.logout() print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] Baostock已登出") def make_prediction(df, predictor): """执行模型推理，仅当天首次调用时运行，添加推理日志""" china_now = get_china_time() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 开始执行模型推理（预测未来{Config['PRED_HORIZON']}个交易日）") # 准备时间戳 last_timestamp = df['timestamps'].max() start_new_range = last_timestamp + pd.Timedelta(days=1) new_timestamps_index = pd.date_range( start=start_new_range, periods=Config["PRED_HORIZON"], freq='D' ) y_timestamp = pd.Series(new_timestamps_index, name='y_timestamp') x_timestamp = df['timestamps'] x_df = df[['open', 'high', 'low', 'close', 'volume', 'amount']] # 推理（禁用梯度计算，节省资源） with torch.no_grad(): begin_time = time.time() close_preds_main, volume_preds_main = predictor.predict( df=x_df, x_timestamp=x_timestamp, y_timestamp=y_timestamp, pred_len=Config["PRED_HORIZON"], T=1.0, top_p=0.95, sample_count=Config["N_PREDICTIONS"], verbose=True ) infer_time = time.time() - begin_time print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 推理完成，耗时{infer_time:.2f}秒") # 波动率预测复用收盘价预测结果（保持原逻辑） close_preds_volatility = close_preds_main return close_preds_main, volume_preds_main, close_preds_volatility def calculate_metrics(hist_df, close_preds_df, v_close_preds_df): """计算上涨概率和波动率放大概率，添加指标计算日志""" print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 开始计算预测指标...") # 上涨概率（最后一个预测日相对于最新收盘价） last_close = hist_df['close'].iloc[-1] final_day_preds = close_preds_df.iloc[-1] upside_prob = (final_day_preds > last_close).mean() # 波动率放大概率（预测波动率vs历史波动率） hist_log_returns = np.log(hist_df['close'] / hist_df['close'].shift(1)) historical_vol = hist_log_returns.iloc[-Config["VOL_WINDOW"]:].std() amplification_count = 0 for col in v_close_preds_df.columns: full_sequence = pd.concat([pd.Series([last_close]), v_close_preds_df[col]]).reset_index(drop=True) pred_log_returns = np.log(full_sequence / full_sequence.shift(1)) predicted_vol = pred_log_returns.std() if predicted_vol > historical_vol: amplification_count += 1 vol_amp_prob = amplification_count / len(v_close_preds_df.columns) # 打印指标日志 print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 指标计算完成：") print(f" - 24个交易日上涨概率：{upside_prob:.2%}") print(f" - 24个交易日波动率放大概率：{vol_amp_prob:.2%}") return upside_prob, vol_amp_prob def create_plot(): china_now = get_china_time() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 开始生成预测图表（适配低版本matplotlib字体）") # 从缓存获取数据（原有逻辑不变） hist_df_for_plot = Config["CACHED_RESULTS"]["hist_df_for_plot"] close_preds = Config["CACHED_RESULTS"]["close_preds"] volume_preds = Config["CACHED_RESULTS"]["volume_preds"] # -------------------------- 新增：创建画布和子图 -------------------------- fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True) # ----------------------------------------------------------------------------- # -------------------------- 修正：低版本matplotlib字体处理 -------------------------- from matplotlib.font_manager import FontProperties font_path = Config["CHINESE_FONT_PATH"] # 检查字体文件是否存在 if os.path.exists(font_path): # 直接通过FontProperties指定字体文件路径（兼容低版本matplotlib） chinese_font = FontProperties(fname=font_path) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 成功加载.ttf字体：{font_path}") else: # 字体文件不存在时的 fallback 逻辑 chinese_font = FontProperties(family='SimHei', size=10) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 字体文件不存在，使用系统默认字体：SimHei") # 全局设置字体（确保坐标轴刻度等默认文本也能显示中文） plt.rcParams["font.family"] = ["sans-serif"] plt.rcParams["font.sans-serif"] = ["WenQuanYi Micro Hei", "SimHei", "Heiti TC"] plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题 # ----------------------------------------------------------------------------- # 绘图时，为所有中文文本显式指定字体（关键） # 1. 价格子图 hist_time = hist_df_for_plot['timestamps'] ax1.plot(hist_time, hist_df_for_plot['close'], color='#00274C', linewidth=1.5) mean_preds = close_preds.mean(axis=1) # 生成预测时间序列（假设预测是在历史最后一个时间之后的24个交易日） last_hist_time = hist_time.max() pred_time = pd.date_range(start=last_hist_time + pd.Timedelta(days=1), periods=Config["PRED_HORIZON"], freq='B') ax1.plot(pred_time, mean_preds, color='#FF6B00', linestyle='-') ax1.fill_between(pred_time, close_preds.min(axis=1), close_preds.max(axis=1), color='#FF6B00', alpha=0.2) # 中文标题/标签指定字体 ax1.set_title(f'{Config["STOCK_CODE"]} 上证指数概率预测（未来{Config["PRED_HORIZON"]}个交易日）', fontsize=16, weight='bold', fontproperties=chinese_font) ax1.set_ylabel('价格（元）', fontsize=12, fontproperties=chinese_font) # 图例指定字体 ax1.legend(['上证指数（后复权）', '预测均价', '预测区间（最小-最大）'], fontsize=10, prop=chinese_font) ax1.grid(True, which='both', linestyle='--', linewidth=0.5) # 2. 成交量子图（同理指定字体） ax2.bar(hist_time, hist_df_for_plot['volume']/1e8, color='#00A86B', width=0.6) ax2.bar(pred_time, volume_preds.mean(axis=1)/1e8, color='#FF6B00', width=0.6) ax2.set_ylabel('成交量（亿手）', fontsize=12, fontproperties=chinese_font) ax2.set_xlabel('日期', fontsize=12, fontproperties=chinese_font) ax2.legend(['历史成交量（亿手）', '预测成交量（亿手）'], fontsize=10, prop=chinese_font) ax2.grid(True, which='both', linestyle='--', linewidth=0.5) # 添加分割线（区分历史和预测数据） separator_time = last_hist_time + pd.Timedelta(hours=12) for ax in [ax1, ax2]: ax.axvline(x=separator_time, color='red', linestyle='--', linewidth=1.5, label='_nolegend_') ax.tick_params(axis='x', rotation=45) # 保存图表 fig.tight_layout() chart_path = Path(Config["CHART_PATH"]) if chart_path.exists(): chart_path.chmod(0o666) # 确保可写权限 fig.savefig(chart_path, dpi=120, bbox_inches='tight') plt.close(fig) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 图表生成完成，保存路径：{chart_path}") def update_html(): """更新HTML页面，复用当天缓存的指标，添加HTML更新日志""" china_now = get_china_time() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 开始更新HTML页面...") # 从缓存获取指标 upside_prob = Config["CACHED_RESULTS"]["upside_prob"] vol_amp_prob = Config["CACHED_RESULTS"]["vol_amp_prob"] now_cn_str = china_now.strftime('%Y-%m-%d %H:%M:%S') upside_prob_str = f'{upside_prob:.1%}' vol_amp_prob_str = f'{vol_amp_prob:.1%}' # ... 原有代码 ... print(f"[DEBUG] 上涨概率：{upside_prob_str}") print(f"[DEBUG] 波动率概率：{vol_amp_prob_str}") # 初始化HTML（不存在则创建基础模板） html_path = Path(Config["HTML_PATH"]) src_html_path = Config["REPO_PATH"] / "templates" / "index.html" if not html_path.exists(): html_path.parent.mkdir(parents=True, exist_ok=True) if src_html_path.exists(): # 复制项目模板 import shutil shutil.copy2(src_html_path, html_path) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 从项目模板复制HTML：{src_html_path} -> {html_path}") else: # 创建基础中文HTML base_html = """ 清华大模型Kronos上证指数预测

清华大学K线大模型Kronos上证指数（sh.000001）概率预测

最后更新时间（中国时间）：未更新

同 步 网 站：火狼工具站

24个交易日上涨概率：--%

24个交易日波动率放大概率：--%

上证指数预测图表
""" with open(html_path, 'w', encoding='utf-8') as f: f.write(base_html) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 在/tmp创建基础HTML：{html_path}") # 更新HTML内容 with open(html_path, 'r', encoding='utf-8') as f: content = f.read() # 替换更新时间 content = re.sub( r'().*?()', lambda m: f'{m.group(1)}{now_cn_str}{m.group(2)}', content ) # 替换上涨概率 content = re.sub( r'().*?()', lambda m: f'{m.group(1)}{upside_prob_str}{m.group(2)}', content ) # 替换波动率放大概率 content = re.sub( r'().*?()', lambda m: f'{m.group(1)}{vol_amp_prob_str}{m.group(2)}', content ) # 写入更新后的HTML with open(html_path, 'w', encoding='utf-8') as f: f.write(content) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] HTML更新完成，路径：{html_path}") # 替换后打印HTML片段 content = re.sub(...) # 上涨概率替换 # 先定义正则表达式 pattern = r'.*?' # 然后在 f-string 中使用 print(f"[DEBUG] 替换前上涨概率片段：{re.search(pattern, content).group(0) if re.search(pattern, content) else '未找到匹配'}") def git_commit_and_push(): """Git提交（仅当Git存在时执行），添加Git操作日志""" china_now = get_china_time() commit_message = f"Auto-update: 上证指数预测（{china_now:%Y-%m-%d 中国时间}）" print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 开始执行Git提交操作，提交信息：{commit_message}") # 检查Git是否安装 try: subprocess.run(['git', '--version'], check=True, capture_output=True, text=True) except (subprocess.CalledProcessError, FileNotFoundError): print(f"[{china_now:%Y-%m-%d %H:%M:%S}] Git未安装或未在PATH中，跳过Git操作") return # 执行Git操作 try: os.chdir(Config["REPO_PATH"]) # 复制图表和HTML到Git跟踪目录（若需要） chart_src = Config["CHART_PATH"] chart_dst = Config["REPO_PATH"] / "prediction_chart.png" html_src = Config["HTML_PATH"] html_dst = Config["REPO_PATH"] / "index.html" if os.path.exists(chart_src): import shutil shutil.copy2(chart_src, chart_dst) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 图表复制到Git目录：{chart_dst}") if os.path.exists(html_src): shutil.copy2(html_src, html_dst) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] HTML复制到Git目录：{html_dst}") # Git add subprocess.run(['git', 'add', 'prediction_chart.png', 'index.html'], check=True, capture_output=True, text=True) # Git commit commit_result = subprocess.run(['git', 'commit', '-m', commit_message], check=True, capture_output=True, text=True) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] Git提交输出：\n{commit_result.stdout}") # Git push push_result = subprocess.run(['git', 'push'], check=True, capture_output=True, text=True) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] Git推送输出：\n{push_result.stdout}") print(f"[{china_now:%Y-%m-%d %H:%M:%S}] Git操作完成") except subprocess.CalledProcessError as e: output = e.stdout if e.stdout else e.stderr if "nothing to commit" in output or "Your branch is up to date" in output: print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 无新内容需要提交或推送") else: print(f"[{china_now:%Y-%m-%d %H:%M:%S}] Git错误：\nSTDOUT: {e.stdout}\nSTDERR: {e.stderr}") except PermissionError as e: print(f"[{china_now:%Y-%m-%d %H:%M:%S}] Git权限错误：{str(e)}，跳过Git操作") def main_task(model): """主任务：控制每日仅执行一次推理，当天复用缓存""" china_now = get_china_time() print(f"\n[{china_now:%Y-%m-%d %H:%M:%S}] " + "="*60) print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 开始执行主任务") # 检查当天是否已完成推理（中国时间） if Config["IS_TODAY_INFERENCED"]: print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 今日（中国时间）已完成推理，直接复用缓存结果") # 复用缓存生成图表和HTML create_plot() update_html() git_commit_and_push() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 主任务完成（复用缓存）") print(f"[{china_now:%Y-%m-%d %H:%M:%S}] " + "="*60 + "\n") return # 当天首次执行：获取数据→推理→缓存结果→生成图表→更新HTML→Git提交 try: # 1. 获取股票数据 df_full = fetch_stock_data() df_for_model = df_full.iloc[:-1] # 排除最后一行避免数据泄漏 # 2. 执行推理 close_preds, volume_preds, v_close_preds = make_prediction(df_for_model, model) # 3. 计算指标 hist_df_for_metrics = df_for_model.tail(Config["VOL_WINDOW"]) upside_prob, vol_amp_prob = calculate_metrics(hist_df_for_metrics, close_preds, v_close_preds) # 4. 缓存结果（当天复用） hist_df_for_plot = df_for_model.tail(Config["VOL_WINDOW"]) # 用于绘图的历史数据 Config["CACHED_RESULTS"] = { "close_preds": close_preds, "volume_preds": volume_preds, "v_close_preds": v_close_preds, "upside_prob": upside_prob, "vol_amp_prob": vol_amp_prob, "hist_df_for_plot": hist_df_for_plot } # 标记当天已完成推理 Config["IS_TODAY_INFERENCED"] = True print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 今日推理结果已缓存，后续调用将复用") # 5. 生成图表 create_plot() # 6. 更新HTML update_html() # 7. Git提交 git_commit_and_push() # 8. 内存回收 del df_full, df_for_model, hist_df_for_metrics gc.collect() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 主任务完成（首次推理）") print(f"[{china_now:%Y-%m-%d %H:%M:%S}] " + "="*60 + "\n") except Exception as e: # 异常时不标记为“已推理”，下次调用重试 Config["IS_TODAY_INFERENCED"] = False print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 主任务执行失败，今日推理标记为未完成") print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 错误信息：{str(e)}") import traceback traceback.print_exc() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] " + "="*60 + "\n") def run_scheduler(model): """定时器：中国时间每天0点触发主任务，其他时间5分钟检查一次""" china_tz = timezone("Asia/Shanghai") print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 定时器启动（中国时间），每天0点执行推理") while True: china_now = get_china_time() # 计算次日0点（中国时间） next_midnight = (china_now + timedelta(days=1)).replace( hour=0, minute=0, second=5, microsecond=0, tzinfo=china_tz ) # 计算等待时间（秒） sleep_seconds = (next_midnight - china_now).total_seconds() # 打印等待日志 print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 当前时间：{china_now:%Y-%m-%d %H:%M:%S}（中国时间）") print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 下次执行时间：{next_midnight:%Y-%m-%d %H:%M:%S}（中国时间）") print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 等待时间：{sleep_seconds:.0f}秒（约{sleep_seconds/3600:.1f}小时）") # 等待到次日0点 time.sleep(sleep_seconds) # 到达0点，执行主任务 try: main_task(model) # 任务完成后，重置“当天已推理”标记（避免跨天复用） Config["IS_TODAY_INFERENCED"] = False except Exception as e: print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 定时器触发任务失败：{str(e)}") import traceback traceback.print_exc() print(f"[{get_china_time():%Y-%m-%d %H:%M:%S}] 5分钟后重试...") time.sleep(300) # 重试间隔5分钟 if __name__ == '__main__': # 初始化：加载模型→执行一次主任务→启动定时器 china_now = get_china_time() print(f"[{china_now:%Y-%m-%d %H:%M:%S}] 程序启动（中国时间）") # 加载模型 loaded_model = load_local_model() # 首次执行主任务（若当天未执行） main_task(loaded_model) # 启动定时器（中国时间每天0点执行） run_scheduler(loaded_model)