""" 批量滚动检测收敛三角形 - 从 pkl 文件读取 OHLCV 数据 每个股票的每个交易日都会计算，输出 DataFrame """ import os import sys import pickle import time import numpy as np import pandas as pd # 让脚本能找到 src/ 下的模块 sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src")) from converging_triangle import ( ConvergingTriangleParams, ConvergingTriangleResult, detect_converging_triangle_batch, ) # ============================================================================ # 【可调参数区】 # ============================================================================ # --- 数据源 --- DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data") # --- 检测参数 --- PARAMS = ConvergingTriangleParams( window=120, pivot_k=15, boundary_n_segments=2, boundary_source="full", upper_slope_max=0.10, lower_slope_min=-0.10, touch_tol=0.10, touch_loss_max=0.10, shrink_ratio=0.8, break_tol=0.001, vol_window=20, vol_k=1.3, false_break_m=5, ) # --- 计算范围 --- # 设为 None 表示计算全部历史；设为具体数字可以只算最近 N 天 RECENT_DAYS = 500 # 默认只算最近 500 天，避免计算时间过长 # --- 输出控制 --- ONLY_VALID = True # True: 只输出识别到三角形的记录; False: 输出所有记录 VERBOSE = True # ============================================================================ # pkl 数据加载 # ============================================================================ class FakeModule: """空壳模块，绕过 model 依赖""" ndarray = np.ndarray def load_pkl(pkl_path: str) -> dict: """加载 pkl 文件，返回字典 {mtx, dtes, tkrs, tkrs_name, ...}""" sys.modules['model'] = FakeModule() sys.modules['model.index_info'] = FakeModule() with open(pkl_path, 'rb') as f: data = pickle.load(f) return data def load_ohlcv_from_pkl(data_dir: str) -> tuple: """ 从 pkl 文件加载 OHLCV 数据 Returns: open_mtx, high_mtx, low_mtx, close_mtx, volume_mtx: shape=(n_stocks, n_days) dates: shape=(n_days,) 真实日期 (如 20050104) tkrs: shape=(n_stocks,) 股票代码 tkrs_name: shape=(n_stocks,) 股票名称 """ open_data = load_pkl(os.path.join(data_dir, "open.pkl")) high_data = load_pkl(os.path.join(data_dir, "high.pkl")) low_data = load_pkl(os.path.join(data_dir, "low.pkl")) close_data = load_pkl(os.path.join(data_dir, "close.pkl")) volume_data = load_pkl(os.path.join(data_dir, "volume.pkl")) # 使用 close 的元数据 dates = close_data["dtes"] tkrs = close_data["tkrs"] tkrs_name = close_data["tkrs_name"] return ( open_data["mtx"], high_data["mtx"], low_data["mtx"], close_data["mtx"], volume_data["mtx"], dates, tkrs, tkrs_name, ) # ============================================================================ # 主流程 # ============================================================================ def main() -> None: start_time = time.time() print("=" * 70) print("Converging Triangle Batch Detection") print("=" * 70) # 1. 加载数据 print("\n[1] Loading OHLCV pkl files...") load_start = time.time() open_mtx, high_mtx, low_mtx, close_mtx, volume_mtx, dates, tkrs, tkrs_name = load_ohlcv_from_pkl(DATA_DIR) load_time = time.time() - load_start n_stocks, n_days = close_mtx.shape print(f" Stocks: {n_stocks}") print(f" Days: {n_days}") print(f" Date range: {dates[0]} ~ {dates[-1]}") print(f" 加载耗时: {load_time:.2f} 秒") # 2. 找有效数据范围（排除全 NaN 的列） any_valid = np.any(~np.isnan(close_mtx), axis=0) valid_day_idx = np.where(any_valid)[0] if len(valid_day_idx) == 0: print("No valid data found!") return last_valid_day = valid_day_idx[-1] print(f" Last valid day: {last_valid_day} (date: {dates[last_valid_day]})") # 3. 计算范围 start_day = PARAMS.window - 1 end_day = last_valid_day # 使用最后有效日，而不是矩阵末尾 if RECENT_DAYS is not None: start_day = max(start_day, end_day - RECENT_DAYS + 1) print(f"\n[2] Detection range: day {start_day} ~ {end_day}") print(f" Window size: {PARAMS.window}") print(f" Total points: {n_stocks} x {end_day - start_day + 1} = {n_stocks * (end_day - start_day + 1)}") # 3. 批量检测 print("\n[3] Running batch detection...") detect_start = time.time() df = detect_converging_triangle_batch( open_mtx=open_mtx, high_mtx=high_mtx, low_mtx=low_mtx, close_mtx=close_mtx, volume_mtx=volume_mtx, params=PARAMS, start_day=start_day, end_day=end_day, only_valid=ONLY_VALID, verbose=VERBOSE, ) detect_time = time.time() - detect_start print(f" 检测耗时: {detect_time:.2f} 秒") # 4. 添加股票代码、名称和真实日期 if len(df) > 0: df["stock_code"] = df["stock_idx"].map(lambda x: tkrs[x] if x < len(tkrs) else "") df["stock_name"] = df["stock_idx"].map(lambda x: tkrs_name[x] if x < len(tkrs_name) else "") df["date"] = df["date_idx"].map(lambda x: dates[x] if x < len(dates) else 0) # 5. 输出结果 print("\n" + "=" * 70) print("Detection Results") print("=" * 70) if ONLY_VALID: print(f"\nTotal valid triangles detected: {len(df)}") else: valid_count = df["is_valid"].sum() print(f"\nTotal records: {len(df)}") print(f"Valid triangles: {valid_count} ({valid_count/len(df)*100:.1f}%)") # 按突破方向统计 if len(df) > 0 and "breakout_dir" in df.columns: breakout_stats = df[df["is_valid"] == True]["breakout_dir"].value_counts() print(f"\nBreakout statistics:") for dir_name, count in breakout_stats.items(): print(f" - {dir_name}: {count}") # 突破强度统计 if len(df) > 0 and "breakout_strength_up" in df.columns: valid_df = df[df["is_valid"] == True] if len(valid_df) > 0: print(f"\nBreakout strength (valid triangles):") print(f" - Up mean: {valid_df['breakout_strength_up'].mean():.4f}, max: {valid_df['breakout_strength_up'].max():.4f}") print(f" - Down mean: {valid_df['breakout_strength_down'].mean():.4f}, max: {valid_df['breakout_strength_down'].max():.4f}") # 6. 保存结果 outputs_dir = os.path.join(os.path.dirname(__file__), "..", "outputs", "converging_triangles") os.makedirs(outputs_dir, exist_ok=True) # 保存完整 CSV (使用 utf-8-sig 支持中文) csv_path = os.path.join(outputs_dir, "all_results.csv") df.to_csv(csv_path, index=False, encoding="utf-8-sig") print(f"\nResults saved to: {csv_path}") # 保存高强度突破记录 (strength > 0.3) if len(df) > 0: strong_up = df[(df["is_valid"] == True) & (df["breakout_strength_up"] > 0.3)] strong_down = df[(df["is_valid"] == True) & (df["breakout_strength_down"] > 0.3)] if len(strong_up) > 0: strong_up_path = os.path.join(outputs_dir, "strong_breakout_up.csv") strong_up.to_csv(strong_up_path, index=False, encoding="utf-8-sig") print(f"Strong up breakouts ({len(strong_up)}): {strong_up_path}") if len(strong_down) > 0: strong_down_path = os.path.join(outputs_dir, "strong_breakout_down.csv") strong_down.to_csv(strong_down_path, index=False, encoding="utf-8-sig") print(f"Strong down breakouts ({len(strong_down)}): {strong_down_path}") # 7. 显示样本 if len(df) > 0: print("\n" + "-" * 70) print("Sample results (first 10):") print("-" * 70) display_cols = [ "stock_code", "date", "is_valid", "breakout_strength_up", "breakout_strength_down", "breakout_dir", "width_ratio" ] display_cols = [c for c in display_cols if c in df.columns] print(df[display_cols].head(10).to_string(index=False)) total_time = time.time() - start_time print("\n" + "=" * 70) print(f"总耗时: {total_time:.2f} 秒 ({total_time/60:.2f} 分钟)") print("=" * 70) if __name__ == "__main__": main()