MoFin/scripts/research/sr_calculator.py

#!/usr/bin/env python3
"""sr_calculator.py — 科学支撑压力计算器（本地数据版，无前视）
对齐 MoFin 算法：
1. 枢轴点系统（calc_support_resistance）：PP/S1/S2/R1/R2 + effective_range
2. 筹码密集区（calc_chip_sr）：640日K线筹码分布 + 2%聚合
数据源：本地 mofin.db（stock_daily），替代腾讯API（可回测、无网络依赖）
"""
import sqlite3, numpy as np
import pandas as pd

class SRCalculator:
    def __init__(self, db="/home/hmo/MoFin/data/mofin.db"):
        self.conn = sqlite3.connect("file:{}?mode=ro".format(db), uri=True)
        self.conn.execute("PRAGMA query_only=ON")
        # 缓存：code -> DataFrame
        self._cache = {}

    def get_bars(self, code):
        """获取个股日线（含前后窗口）"""
        if code in self._cache:
            return self._cache[code]
        df = pd.read_sql(
            "SELECT date, open, close, high, low, volume FROM stock_daily WHERE code=? ORDER BY date",
            self.conn, params=(code,))
        df["date"] = df["date"].astype(str)
        self._cache[code] = df
        return df

    def pivot_points(self, code, date_idx, lookback=10):
        """枢轴点系统：用最近 lookback 日（含当日）的 H/L/C 计算
        返回: {pp, s1, s2, r1, r2, effective_range}
        """
        df = self.get_bars(code)
        if date_idx < 0 or date_idx >= len(df):
            return None
        # 用当日 + 前 lookback 日窗口
        win = df.iloc[max(0, date_idx-lookback+1):date_idx+1]
        if len(win) < 3:
            return None
        h = win["high"].max()
        l = win["low"].min()
        c = win["close"].iloc[-1]
        if not all([h, l, c]) or h <= 0 or l <= 0 or c <= 0:
            return None
        # 有效区间 = max(窗口波幅, 价格×5%)
        daily_range = win["high"].iloc[-1] - win["low"].iloc[-1]
        multi_range = h - l
        min_range = c * 0.05
        effective_range = max(daily_range, multi_range, min_range)
        # 高位/低位扩大
        if h > l:
            trend_pos = (c - l) / (h - l)
            if trend_pos > 0.8 or trend_pos < 0.2:
                effective_range = max(effective_range, c * 0.08)
        # 枢轴点
        pp = (h + l + c) / 3
        s1 = 2 * pp - h
        s2 = pp - effective_range
        r1 = 2 * pp - l
        r2 = pp + effective_range
        return {
            "pp": pp, "s1": s1, "s2": s2, "r1": r1, "r2": r2,
            "effective_range": effective_range,
            "multi_high": h, "multi_low": l,
        }

    def chip_sr(self, code, date_idx, lookback=640):
        """筹码密集区：用 date_idx 之前 lookback 日构建筹码分布
        返回: {chip_ss, chip_sr} 或 None
        """
        df = self.get_bars(code)
        if date_idx < 0 or date_idx >= len(df):
            return None
        price = df["close"].iloc[date_idx]
        if price <= 0:
            return None
        win = df.iloc[max(0, date_idx-lookback):date_idx+1]
        if len(win) < 30:
            return None
        # 构建筹码分布（对齐 MoFin：OHLC 区间均匀分配 + 衰减）
        chip_dist = {}
        decay = 0.97
        n = len(win)
        for k, row in enumerate(win.itertuples()):
            high, low, volume = row.high, row.low, row.volume
            if high <= low or volume <= 0:
                continue
            step = max(round((high - low) / 5, 2), 0.01)
            level = round(low, 2)
            vol_per_level = volume / max(int((high - low) / step) + 1, 1)
            while level <= high:
                chip_dist[level] = chip_dist.get(level, 0) + vol_per_level
                level = round(level + step, 2)
        if not chip_dist:
            return None
        # 2% 区间聚合
        step = max(round(price * 0.02, 2), 1.0)
        bins = {}
        for p, v in chip_dist.items():
            k = round(p / step) * step
            bins[k] = bins.get(k, 0) + v
        sb = sorted(bins.items())
        below = [(p, v) for p, v in sb if p < price]
        above = [(p, v) for p, v in sb if p >= price]
        if not below or not above:
            return None
        chip_ss = max(below, key=lambda x: x[1])[0]
        chip_sr = max(above, key=lambda x: x[1])[0]
        return {"chip_ss": chip_ss, "chip_sr": chip_sr}

    def sr_full(self, code, date_idx):
        """综合支撑压力：枢轴点 + 筹码密集区 + 共振判断
        返回支撑/压力位 + 强弱标签
        """
        pv = self.pivot_points(code, date_idx)
        chip = self.chip_sr(code, date_idx)
        result = {"code": code, "pivot": pv, "chip": chip}
        if pv:
            # 支撑候选：S1/S2/筹码支撑
            cands_s = [("pivot_s1", pv["s1"]), ("pivot_s2", pv["s2"])]
            if chip:
                cands_s.append(("chip_ss", chip["chip_ss"]))
            # 压力候选
            cands_r = [("pivot_r1", pv["r1"]), ("pivot_r2", pv["r2"])]
            if chip:
                cands_r.append(("chip_sr", chip["chip_sr"]))
            result["cands_s"] = cands_s
            result["cands_r"] = cands_r
        return result

    def close(self):
        self.conn.close()

if __name__ == "__main__":
    # 自测：茅台 600519 某日
    sr = SRCalculator()
    df = sr.get_bars("600519")
    print("600519 行数:", len(df))
    # 取最后第 10 天（留出未来模拟空间）
    idx = len(df) - 10
    date = df["date"].iloc[idx]
    price = df["close"].iloc[idx]
    print("测试日:", date, "价格:", price)
    pv = sr.pivot_points("600519", idx)
    print("枢轴点:", {k: round(v, 2) for k, v in pv.items() if isinstance(v, float)})
    chip = sr.chip_sr("600519", idx)
    print("筹码:", {k: round(v, 2) for k, v in chip.items()} if chip else None)
    sr.close()