quant-trading/research/lab/panel.py

"""
Build research Panels (see engine.Panel) from the fetched raw data.

    from research.lab import panel
    us = panel.load_us()          # S&P 500 PIT, 2004 →
    cn = panel.load_cn("hs300")   # CSI 300 PIT, 2007 →  ("zz800" = CSI 300 + 500)

Panels are cached in data/research/*.panel.pkl; delete them after refetching.
"""

from __future__ import annotations

import json
import os
import pickle

import numpy as np
import pandas as pd

from research.lab.engine import Costs, Panel
from research.lab.tiingo_merge import apply_sp500

RAW_DIR = os.path.join("data", "research")
US_ETFS = {"SPY", "RSP", "QQQ", "SHY", "BIL", "IEF", "TLT", "GLD",
           "XLB", "XLE", "XLF", "XLI", "XLK", "XLP", "XLU", "XLV", "XLY"}


def _cached(name: str, build):
    path = os.path.join(RAW_DIR, f"{name}.panel.pkl")
    if os.path.exists(path):
        with open(path, "rb") as fh:
            return pickle.load(fh)
    p = build()
    with open(path, "wb") as fh:
        pickle.dump(p, fh)
    return p


# ---------------------------------------------------------------------------
# US
# ---------------------------------------------------------------------------

def _us_membership(dates: pd.DatetimeIndex, tickers: list[str]) -> pd.DataFrame:
    with open(os.path.join("data", "sp500_history.json")) as fh:
        intervals = json.load(fh)["intervals"]
    mask = pd.DataFrame(False, index=dates, columns=tickers)
    for t in tickers:
        for start, end in intervals.get(t, []):
            s = pd.Timestamp(start) if start else dates[0]
            e = pd.Timestamp(end) if end else dates[-1] + pd.Timedelta(days=1)
            mask.loc[(dates >= s) & (dates < e), t] = True
    return mask


def _build_us() -> Panel:
    with open(os.path.join(RAW_DIR, "us_ohlcv.pkl"), "rb") as fh:
        raw = pickle.load(fh)
    raw = apply_sp500(raw)
    close = raw["close"]
    # keep rows where most listed tickers printed (drops stray holiday rows)
    alive = close.notna().sum(axis=1)
    close = close[alive >= 0.5 * alive.rolling(20, min_periods=1).max()]
    dates = close.index
    etf = {k: raw[k].reindex(dates)[[c for c in raw[k].columns if c in US_ETFS]]
           for k in ("close",)}["close"]
    stocks = sorted(c for c in close.columns if c not in US_ETFS)
    f = {k: raw[k].reindex(index=dates, columns=stocks) for k in ("open", "high", "low", "close", "volume")}
    # Yahoo sometimes prints zero/negative adjusted opens for old bars
    for k in ("open", "high", "low"):
        f[k] = f[k].where(f[k] > 0)
    has_px = f["close"].notna()
    member = _us_membership(dates, stocks) & has_px
    cash_px = etf["BIL"].combine_first(etf["SHY"]) if "BIL" in etf else etf["SHY"]
    bil = etf["BIL"].pct_change(fill_method=None)
    shy = etf["SHY"].pct_change(fill_method=None)
    cash_ret = bil.where(bil.notna(), shy).fillna(0.0)
    del cash_px
    return Panel(
        market="us", open=f["open"], high=f["high"], low=f["low"], close=f["close"],
        volume=f["volume"], member=member, can_buy=has_px, can_sell=has_px,
        cash_ret=cash_ret, benchmark=etf["SPY"],
        costs=Costs(buy=0.001, sell=0.001, fixed_fee=0.0, capital=10_000.0),
        extra={"etf": etf, "sector": _us_sectors(stocks)},
    )


def _us_sectors(stocks: list[str]) -> pd.Series:
    """GICS sector per ticker; current members from data/us_sectors.csv, the
    rest from data/research/us_sectors_extra.csv when it has been fetched."""
    sec = pd.read_csv(os.path.join("data", "us_sectors.csv")).set_index("ticker")["GICS Sector"]
    extra = os.path.join(RAW_DIR, "us_sectors_extra.csv")
    if os.path.exists(extra):
        sec = sec.combine_first(pd.read_csv(extra).set_index("ticker")["sector"])
    return sec.reindex(stocks)


def _tiingo_sig() -> list:
    d = os.path.join(RAW_DIR, "tiingo")
    if not os.path.isdir(d):
        return [0, 0]
    files = [f for f in os.listdir(d) if f.endswith(".csv")]
    mtime = max((os.path.getmtime(os.path.join(d, f)) for f in files), default=0)
    return [len(files), round(mtime, 3)]


def load_us() -> Panel:
    """Cached US panel, rebuilt whenever the Tiingo overlay changes."""
    path = os.path.join(RAW_DIR, "us.panel.pkl")
    key = os.path.join(RAW_DIR, "us.panel.key")
    sig = _tiingo_sig()
    if os.path.exists(path) and os.path.exists(key):
        with open(key) as fh:
            if json.load(fh) == sig:
                with open(path, "rb") as fh:
                    return pickle.load(fh)
    p = _build_us()
    with open(path, "wb") as fh:
        pickle.dump(p, fh)
    with open(key, "w") as fh:
        json.dump(sig, fh)
    return p


# ---------------------------------------------------------------------------
# CN
# ---------------------------------------------------------------------------

def _limit_pct(codes: list[str], dates: pd.DatetimeIndex, isst: pd.DataFrame) -> pd.DataFrame:
    """Daily price-limit fraction per stock (10% main board, 20% ChiNext after
    2020-08-24 and STAR, 5% ST)."""
    lim = pd.DataFrame(0.10, index=dates, columns=codes)
    star = [c for c in codes if c.split(".")[1].startswith("688")]
    chinext = [c for c in codes if c.split(".")[1].startswith("30")]
    lim[star] = 0.20
    lim.loc[dates >= "2020-08-24", chinext] = 0.20
    st = isst.reindex(index=dates, columns=codes).fillna(0).astype(bool)
    lim = lim.mask(st & (lim < 0.2), 0.05)
    return lim


def _cn_membership(members: dict[str, list[str]], dates: pd.DatetimeIndex,
                   codes: list[str]) -> pd.DataFrame:
    snap = sorted(members.items())
    mask = pd.DataFrame(False, index=dates, columns=codes)
    col = {c: i for i, c in enumerate(codes)}
    arr = mask.values
    for i, (d, lst) in enumerate(snap):
        s = pd.Timestamp(d)
        e = pd.Timestamp(snap[i + 1][0]) if i + 1 < len(snap) else dates[-1] + pd.Timedelta(days=1)
        rows = (dates >= s) & (dates < e)
        cols = [col[c] for c in lst if c in col]
        arr[np.ix_(rows, cols)] = True
    return pd.DataFrame(arr, index=dates, columns=codes)


def _build_cn(index_set: str) -> Panel:
    with open(os.path.join(RAW_DIR, "cn_ohlcv.pkl"), "rb") as fh:
        raw = pickle.load(fh)
    F = raw["fields"]
    # Codes an earlier fetcher filled from Tencent's 后复权 series are dropped:
    # it damps returns for long dividend histories (limit-down shows -7%).
    bad = {c for c, src in raw.get("source", {}).items() if src == "tencent"}
    close = F["close"]
    dates = close.index
    codes = [c for c in close.columns if c not in bad]
    f = {k: v.reindex(index=dates, columns=codes) for k, v in F.items()}
    # Yahoo bars lack ST flags and valuation ratios; index members
    # are rarely ST (the index drops them at the next review).
    if "isst" not in f:
        f["isst"] = pd.DataFrame(0.0, index=dates, columns=codes)
    traded = (f["tradestatus"] == 1) & (f["volume"] > 0) & f["close"].notna()
    if index_set == "hs300":
        member = _cn_membership(raw["members"]["hs300"], dates, codes)
    elif index_set == "zz500":
        member = _cn_membership(raw["members"]["zz500"], dates, codes)
    else:  # zz800
        member = (_cn_membership(raw["members"]["hs300"], dates, codes)
                  | _cn_membership(raw["members"]["zz500"], dates, codes))
    st = f["isst"].fillna(0).astype(bool)
    member = member & traded.rolling(5, min_periods=1).max().astype(bool) & ~st
    lim = _limit_pct(codes, dates, f["isst"])
    open_ret = f["open"] / f["preclose"] - 1
    can_buy = traded & ~(open_ret >= lim - 0.002)
    can_sell = traded & ~(open_ret <= -lim + 0.002)
    # halted days carry the last price, so returns across a halt are zero
    for k in ("open", "high", "low", "close"):
        f[k] = f[k].where(traded)
    stamp = pd.Series(0.001, index=dates)
    stamp[dates >= "2023-08-28"] = 0.0005
    idx = raw["index"].reindex(dates).ffill()
    return Panel(
        market=f"cn_{index_set}", open=f["open"], high=f["high"], low=f["low"], close=f["close"],
        volume=f["volume"].where(traded), member=member, can_buy=can_buy, can_sell=can_sell,
        cash_ret=pd.Series(0.02 / 252, index=dates), benchmark=idx["sh.000300"],
        costs=Costs(buy=0.0008, sell=0.0008, sell_tax=stamp, fixed_fee=0.0, capital=70_000.0),
        extra=_cn_extra(f, traded, lim, idx),
    )


def _cn_extra(f: dict, traded: pd.DataFrame, lim: pd.DataFrame, idx: pd.DataFrame) -> dict:
    extra = {"limit": lim, "index": idx, "source": "baostock" if "pe" in f else "yahoo"}
    if "turn" in f:
        extra["turn"] = f["turn"].where(traded)
    if "amount" in f:
        extra["amount"] = f["amount"].where(traded)
        if "turn" in f:
            extra["float_cap"] = (f["amount"] / (f["turn"] / 100)).where(traded & (f["turn"] > 0))
    for k in ("pe", "pb", "ps", "pcf"):
        if k in f:
            extra[k] = f[k]
    return extra


def load_cn(index_set: str = "hs300") -> Panel:
    return _cached(f"cn_{index_set}", lambda: _build_cn(index_set))


def coverage_report(p: Panel) -> pd.DataFrame:
    """Per year: PIT members with price data vs expected index size."""
    m = p.member.resample("YE").last()
    return pd.DataFrame({"members_with_data": m.sum(axis=1)})