""" 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 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) 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 load_us() -> Panel: return _cached("us", _build_us) # --------------------------------------------------------------------------- # 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)})