quant-trading/research/lab/panel_mid.py

"""
Point-in-time US mid- and small-cap research Panels (see engine.Panel), built
from iShares IJH / IJR month-end holdings and Yahoo bars fetched by
research/lab/fetch_us_mid.py.

    from research.lab.panel_mid import load_us_mid
    p = load_us_mid("mid")        # S&P MidCap 400 members, 2006-09-29 →
    p = load_us_mid("small")      # S&P SmallCap 600
    p = load_us_mid("midsmall")   # the union (S&P 400 + 600)

``member`` is True when the ticker is in the fund's latest month-end snapshot
on or before that date AND Yahoo has a close that day, so index changes show
up at the next month-end (up to a month late; no snapshots exist for
2017-01..2017-06, so the Dec-2016 list stands until Jul-2017).

Same-security check. The holdings files carry each stock's as-traded price;
Yahoo's close with later splits undone reproduces it to ~1e-8. A snapshot
row counts only when the two agree (``price_check``), which rejects symbols
Yahoo now serves for another company (Barnes Group's "B" is Barrick,
CoreSite's "COR" is Cencora). Most delisted names are missing on Yahoo:
``survivorship_coverage`` measures how much of each snapshot has prices, and
``member_vs_etf`` compares the covered members with the ETF.

Cleaning: Yahoo's stray ticks (``bad_prints``) are removed, and a series that
turns into another security mid-membership (``_splices``) is cut there.
can_buy = can_sell = has a price. Costs are 20 bp a side (mid/small spreads
are wider than the S&P 500's). Benchmark: IJH (mid), IJR (small), their
daily equal blend (midsmall); extra["etf"] holds IJH, IJR, MDY, IWM, SPY,
QQQ, BIL, SHY closes, extra["index_weight"] each name's weight in its fund
at the latest snapshot.

Caches: data/research/us_mid_bars.pkl (all bars) and
data/research/us_{mid,small,midsmall}.panel.pkl; delete them after refetching.
"""

from __future__ import annotations

import os
import pickle

import numpy as np
import pandas as pd

from research.lab.engine import Costs, Panel
from research.lab.fetch_us_mid import ETFS, load_bars, load_holdings

RAW_DIR = os.path.join("data", "research")
FUNDS = {"mid": ("IJH",), "small": ("IJR",), "midsmall": ("IJH", "IJR")}
START = "2005-01-01"
PRICE_TOL = 0.02          # |log(iShares price / Yahoo as-traded close)|; true matches are ~1e-8
JUMP = np.log(3)          # a one-day move this large is checked (bad print, splice)
STABLE_TOL = 0.005        # ... or an offset (merger ratio, spin-off) this steady over 3 snapshots
MAX_OFFSET = 0.5          # largest steady offset accepted
SECTOR_ALIASES = {"Communication": "Communication Services",
                  "Telecommunications": "Communication Services",
                  "Telecommunication Services": "Communication Services"}


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, protocol=pickle.HIGHEST_PROTOCOL)
    return p


def bad_prints(close: np.ndarray, volume: np.ndarray, window: int = 5) -> np.ndarray:
    """Bars more than JUMP (3x) away from the last good close that either did
    not trade (zero volume, or under 1% of the trailing 60-day median) or
    come back within 1.5x of it inside ``window`` days: Yahoo's stray ticks."""
    n = len(close)
    bad = np.zeros(n, bool)
    medv = pd.Series(volume).rolling(60, min_periods=10).median().shift(1).to_numpy()
    fin = np.isfinite(close) & (close > 0)
    last = np.nan
    for i in range(n):
        if not fin[i]:
            continue
        if last == last and abs(np.log(close[i] / last)) > JUMP:
            thin = volume[i] == 0 or (medv[i] == medv[i] and volume[i] < 0.01 * medv[i])
            nxt = close[i + 1:i + 1 + window]
            nxt = nxt[np.isfinite(nxt) & (nxt > 0)]
            if thin or (len(nxt) and np.any(np.abs(np.log(nxt / last)) < np.log(1.5))):
                bad[i] = True
                continue
        last = close[i]
    return bad


def _bars_all() -> dict:
    """Wide adjusted OHLCV + as-traded close for every ever-member and ETF,
    on SPY's trading calendar from 2005 (cached: rebuilding reads ~3k CSVs).
    ``bad_prints`` are removed from OHLCV (``bad`` marks them); ``raw_close``
    is left as Yahoo serves it for the price check."""
    path = os.path.join(RAW_DIR, "us_mid_bars.pkl")
    if os.path.exists(path):
        with open(path, "rb") as fh:
            return pickle.load(fh)
    hold = load_holdings()
    tickers = sorted(set(hold["ticker"]) | set(ETFS))
    spy = load_bars("SPY")
    dates = spy.index[spy.index >= START]
    fields = ("open", "high", "low", "close", "volume", "raw_close")
    arr = {k: np.full((len(dates), len(tickers)), np.nan) for k in fields}
    have = []
    for j, t in enumerate(tickers):
        b = load_bars(t)
        if b is None:
            continue
        b = b.reindex(dates)
        for k in fields:
            arr[k][:, j] = b[k].to_numpy(float)
        have.append(t)
    out = {k: pd.DataFrame(v, index=dates, columns=tickers)[have] for k, v in arr.items()}
    bad = np.column_stack([bad_prints(out["close"][t].to_numpy(), out["volume"][t].to_numpy())
                           for t in have])
    out["bad"] = pd.DataFrame(bad, index=dates, columns=have)
    for k in ("open", "high", "low", "close", "volume"):
        out[k] = out[k].mask(out["bad"])
    with open(path, "wb") as fh:
        pickle.dump(out, fh, protocol=pickle.HIGHEST_PROTOCOL)
    return out


def price_check(hold: pd.DataFrame, raw_close: pd.DataFrame) -> pd.DataFrame:
    """Per holding row: log of the iShares share price over Yahoo's as-traded
    close on the snapshot date (last close within 5 days). A row matches when
    they agree within PRICE_TOL, or when the ratio holds within STABLE_TOL
    over three consecutive snapshots of the ticker: Yahoo scales some
    histories by a constant (a spin-off: VSH before 2010 is off by 10%; a
    merger's exchange ratio: DCOM before 2021 is old Dime x 0.648). The
    result is then the majority of each snapshot and its two neighbours,
    which drops a lone coincidental match of another company and keeps a
    lone event-month glitch. ``has_px``: Yahoo has a close."""
    rc = raw_close.ffill(limit=5)
    h = hold[["fund", "date", "ticker", "price", "weight"]].sort_values(["fund", "ticker", "date"])
    rows = rc.index.get_indexer(h["date"], method="ffill")
    cols = rc.columns.get_indexer(h["ticker"])
    ok = (rows >= 0) & (cols >= 0)
    yc = np.full(len(h), np.nan)
    yc[ok] = rc.to_numpy()[rows[ok], cols[ok]]
    h["yahoo_close"] = yc
    h["has_px"] = np.isfinite(yc) & (yc > 0)
    with np.errstate(divide="ignore", invalid="ignore"):
        lr = np.log(h["price"] / h["yahoo_close"])
    h["log_ratio"] = lr
    key = [h["fund"], h["ticker"]]
    g = lr.groupby(key)
    near_prev = (lr - g.shift(1)).abs() <= STABLE_TOL
    near_next = (lr - g.shift(-1)).abs() <= STABLE_TOL
    # an offset held over three snapshots in a row (two unrelated prices rarely
    # do that), and the run's first and last snapshot
    core = near_prev & near_next & (lr.abs() <= MAX_OFFSET)
    cg = core.groupby(key)
    steady = (core | (near_next & cg.shift(-1, fill_value=False).astype(bool))
              | (near_prev & cg.shift(1, fill_value=False).astype(bool)))
    raw = h["has_px"] & ((lr.abs() <= PRICE_TOL) | steady | h["price"].isna())
    hp = h["has_px"].groupby(key)
    has_prev = hp.shift(1, fill_value=False).astype(bool)
    has_next = hp.shift(-1, fill_value=False).astype(bool)
    rg = raw.groupby(key)
    votes = (raw.astype(int) + (rg.shift(1, fill_value=False).astype(bool) & has_prev).astype(int)
             + (rg.shift(-1, fill_value=False).astype(bool) & has_next).astype(int))
    both = has_prev & has_next
    h["match_raw"] = raw
    h["match"] = h["has_px"] & np.where(both, votes >= 2, raw)
    return h.sort_index()


def _splices(chk: pd.DataFrame, bars: dict) -> pd.DataFrame:
    """Dates whose Yahoo prices belong to another security spliced onto the
    series (Intermec's "IN" turns into a penny stock in Nov 2012): the
    ticker matches at one snapshot, has a price that does not match at the
    fund's next snapshot, and the series jumps (or prints badly) in between.
    Prices are dropped from that day until the ticker matches again."""
    close = bars["close"]
    jump = (np.log(close.ffill()).diff().abs() > JUMP) | bars["bad"]
    drop = pd.DataFrame(False, index=close.index, columns=close.columns)
    matched = chk[chk["match"]].groupby("ticker")["date"].apply(lambda d: np.sort(d.unique()))
    for fund, gf in chk.groupby("fund"):
        snaps = np.sort(gf["date"].unique())
        nxt = dict(zip(snaps[:-1], snaps[1:]))
        for t, g in gf.sort_values("date").groupby("ticker"):
            if t not in close.columns:
                continue
            d, m, px = g["date"].to_numpy(), g["match"].to_numpy(), g["has_px"].to_numpy()
            for i in range(len(g) - 1):
                if not (m[i] and px[i + 1] and not m[i + 1] and nxt.get(d[i]) == d[i + 1]):
                    continue
                win = jump.loc[(jump.index > d[i]) & (jump.index <= d[i + 1]), t]
                if not win.any():
                    continue
                start = win.idxmax()
                later = matched.get(t, np.array([], dtype="datetime64[ns]"))
                later = later[later > d[i + 1]]
                end = later[0] if len(later) else close.index[-1] + pd.Timedelta(days=1)
                drop.loc[(drop.index >= start) & (drop.index < end), t] = True
    return drop


def _membership(chk: pd.DataFrame, dates: pd.DatetimeIndex, tickers: list[str]) -> pd.DataFrame:
    """In the fund's latest snapshot on or before each date, with a matching price."""
    col = {t: i for i, t in enumerate(tickers)}
    arr = np.zeros((len(dates), len(tickers)), bool)
    for fund, g in chk.groupby("fund"):
        snaps = sorted(g["date"].unique())
        for i, s in enumerate(snaps):
            e = snaps[i + 1] if i + 1 < len(snaps) else dates[-1] + pd.Timedelta(days=1)
            rows = (dates >= s) & (dates < e)
            names = g.loc[(g["date"] == s) & g["match"], "ticker"]
            cols = [col[t] for t in names if t in col]
            arr[np.ix_(rows, cols)] = True
    return pd.DataFrame(arr, index=dates, columns=tickers)


def _index_weight(chk: pd.DataFrame, dates, tickers) -> pd.DataFrame:
    """Weight in its fund at the latest snapshot (held constant until the next)."""
    w = chk.pivot_table(index="date", columns="ticker", values="weight", aggfunc="sum")
    w = w.reindex(columns=tickers)
    return w.reindex(dates.union(w.index)).ffill().reindex(dates).fillna(0.0).astype("float32")


def _build(which: str) -> Panel:
    funds = FUNDS[which]
    bars = _bars_all()
    hold = load_holdings()
    hold = hold[hold["fund"].isin(funds)]
    chk = price_check(hold, bars["raw_close"])
    close_all = bars["close"]
    dates = close_all.index
    tickers = sorted(set(chk.loc[chk["match"], "ticker"]) & set(close_all.columns) - set(ETFS))
    spliced = _splices(chk, bars)[tickers]
    f = {k: bars[k][tickers].mask(spliced) for k in ("open", "high", "low", "close", "volume")}
    for k in ("open", "high", "low"):
        f[k] = f[k].where(f[k] > 0)
    f["close"] = f["close"].where(f["close"] > 0)
    has_px = f["close"].notna()
    member = _membership(chk, dates, tickers) & has_px
    etf = close_all[[e for e in ETFS if e in close_all.columns]]
    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)
    if which == "midsmall":      # S&P 400 + 600 → equal blend of the two funds
        r = etf[["IJH", "IJR"]].pct_change(fill_method=None).mean(axis=1).fillna(0.0)
        bench = (1 + r).cumprod()
    else:
        bench = etf[funds[0]]
    # the ticker's latest GICS sector (as panel.py, today's sector for all dates)
    sector = (hold[~hold["sector"].isin(["-", "", "Other"])].sort_values("date")
              .groupby("ticker")["sector"].last().replace(SECTOR_ALIASES).reindex(tickers))
    return Panel(
        market=f"us_{which}", 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=bench,
        costs=Costs(buy=0.002, sell=0.002, fixed_fee=0.0, capital=10_000.0),
        extra={"etf": etf, "sector": sector,
               "index_weight": _index_weight(chk, dates, tickers)},
    )


def load_us_mid(which: str = "mid") -> Panel:
    if which not in FUNDS:
        raise ValueError(f"which must be one of {list(FUNDS)}")
    return _cached(f"us_{which}", lambda: _build(which))


# ---------------------------------------------------------------------------
# diagnostics
# ---------------------------------------------------------------------------

def survivorship_coverage(which: str = "midsmall") -> pd.DataFrame:
    """Per fund and snapshot: members, and the share (by count and by index
    weight) with a Yahoo price for the same security on that date."""
    bars = _bars_all()
    hold = load_holdings()
    chk = price_check(hold[hold["fund"].isin(FUNDS[which])], bars["raw_close"])
    chk["w_ok"] = chk["weight"] * chk["match"]
    chk["w_px"] = chk["weight"] * chk["has_px"]
    g = chk.groupby(["fund", "date"])
    return pd.DataFrame({
        "n": g.size(), "n_ok": g["match"].sum(), "count_cov": g["match"].mean(),
        "weight_cov": g["w_ok"].sum() / g["weight"].sum(),
        "count_any_px": g["has_px"].mean(),
        "weight_any_px": g["w_px"].sum() / g["weight"].sum(),
    })


def member_vs_etf(p: Panel) -> pd.DataFrame:
    """Calendar-year return of the equal-weight member portfolio (daily
    rebalanced), the index-weighted covered members, and the ETF."""
    r = p.close.pct_change(fill_method=None)
    prev = p.member.shift(1, fill_value=False)
    ew = r.where(prev).mean(axis=1)
    w = p.extra["index_weight"].shift(1).where(prev).astype(float)
    cw = (r.where(prev) * w).sum(axis=1) / w.where(r.where(prev).notna()).sum(axis=1)
    bench = p.benchmark.pct_change(fill_method=None)
    df = pd.DataFrame({"ew_members": ew, "cw_members": cw, "etf": bench}).loc[prev.any(axis=1)]
    yr = (1 + df.fillna(0.0)).groupby(df.index.year).prod() - 1
    yr["ew_minus_etf"] = yr["ew_members"] - yr["etf"]
    yr["cw_minus_etf"] = yr["cw_members"] - yr["etf"]
    return yr