quant-trading/research/lab/lead_lag_small.py

"""
Lead-lag and supply-chain signals on S&P 400 + 600 stocks, where the
literature places these effects (low attention, large customers leading small
suppliers). Companion to research/lab/lead_lag.py (S&P 500).

    PYTHONPATH=. uv run python -m research.lab.fetch_supply_links --universe mid   # 10-K links, ~1 h
    PYTHONPATH=. uv run --with numba python -m research.lab.lead_lag_small

Universe: PIT S&P 400 + 600 members (research/lab/panel_mid.py, iShares
holdings; covers 66% of members in 2019 and 83% in 2023; the missing names are
mostly delisted ones, so long-only returns carry survivorship bias) as the
stocks traded, and the S&P 1500 (those plus PIT S&P 500 members) as the pool
of linked names. Main sample 2019-01..2026-09, halves 2019-22 / 2023-26,
2017-26 for reference. Every signal is fixed before the results are seen and
tested at both horizons h = 21 and 63 trading days (signal window = holding
period = h), because the AI-chain stages lagged each other by 2-6 months.

Signals for small stock i, known at the close of day t:
  peer_mom / peer_gap    the 10 S&P 1500 members whose daily residual returns
                         (S&P 1500 equal-weight beta removed) correlated most
                         with i over the last 252 days, re-estimated each
                         month end: their mean log return over the last h days
                         (mom), and that minus i's own (gap).
  bigpeer_mom            the same with the 5 most correlated S&P 500 members
                         only (large firms leading small ones, Hou 2007).
  cust_mom / cust_gap    S&P 1500 members named as customers in i's 10-Ks
                         (data/research/mid_10k_links.csv plus S&P 500-run
                         links whose filer is a mid/small member), valid for
                         15 months from the filing day.
Event test: after an S&P 500 member's earnings day (8-K item 2.02), the
abnormal returns of the small stocks that have it among their 5 big peers, on
the day and over days 1-21 and 1-63, regressed on its day-0 abnormal return.

Measures per signal and horizon (samples every 21 days; Newey-West t with
h/21 - 1 lags):
  * Fama-MacBeth slope of the next-h open-to-open return rank on the signal
    rank, controlling for the stock's 1-month return and 12-1 momentum ranks.
  * top-minus-bottom quintile return, and the top quintile minus the
    equal-weight small universe after 20 bp a side on its measured turnover.
  * a top-30 sleeve (engine.simulate, next-open fills, 20 bp, buffer 2,
    rebalanced every h days, phase-averaged): alpha against the IJH/IJR blend
    and against QQQ.
The S&P 500 members get the same peer_mom / peer_gap / cust (S&P 500 10-K
links) FM rows over the same years for comparison.
"""

from __future__ import annotations

import json
import os
import warnings

import numpy as np
import pandas as pd

from research.lab.factors import _rolling_beta

warnings.filterwarnings("ignore")
DIR = os.path.join("data", "research")
RESULTS = os.path.join(DIR, "results")
END = "2026-09-24"
SPANS = {"2019-26": ("2019-01-01", END), "2019-22": ("2019-01-01", "2022-12-31"),
         "2023-26": ("2023-01-01", END), "2017-26": ("2017-01-01", END)}
HS = (21, 63)


# ---------------------------------------------------------------------------
# S&P 1500 frame
# ---------------------------------------------------------------------------

def frame():
    from research.lab.factor_scan import load
    from research.lab.panel_mid import load_us_mid
    pb, _ = load("us")
    ps = load_us_mid("midsmall")
    dates = pb.dates.intersection(ps.dates)
    dates = dates[dates <= pd.Timestamp(END)]
    cols = sorted(set(pb.close.columns) | set(ps.close.columns))

    def comb(a, b):
        return a.reindex(index=dates, columns=cols).combine_first(b.reindex(index=dates, columns=cols))
    close = comb(pb.close, ps.close)
    opn = comb(pb.open, ps.open)
    mb = pb.member.reindex(index=dates, columns=cols).fillna(False).astype(bool)
    ms = ps.member.reindex(index=dates, columns=cols).fillna(False).astype(bool) & ~mb
    return {"pb": pb, "ps": ps, "dates": dates, "cols": cols, "close": close, "open": opn,
            "mb": mb, "ms": ms, "mall": mb | ms}


def residuals(F) -> pd.DataFrame:
    r = F["close"].pct_change(fill_method=None)
    m = r.where(F["mall"]).mean(axis=1).fillna(0.0)
    beta = _rolling_beta(r, m, 252, 126).shift(1)
    F["ret"], F["mkt"] = r, m
    return r - beta.mul(m, axis=0)


def month_ends(dates) -> list[pd.Timestamp]:
    s = pd.Series(dates, index=dates)
    return list(s.groupby(dates.to_period("M")).max())


def peer_links(F, e: pd.DataFrame, focal: pd.DataFrame, pool: pd.DataFrame, k: int, win: int = 252,
               start: str = "2015-06-01") -> dict:
    """{month_end: {focal col index: peer col indices}} by residual correlation."""
    X = e.where(F["mall"]).to_numpy()
    out = {}
    fo, po = focal.to_numpy(), pool.to_numpy()
    for me in month_ends(F["dates"]):
        if me < pd.Timestamp(start):
            continue
        i1 = F["dates"].get_loc(me) + 1
        blk = X[i1 - win:i1]
        ok = np.isfinite(blk).sum(axis=0) >= int(win * 0.8)
        fi = np.where(fo[i1 - 1] & ok)[0]
        pi = np.where(po[i1 - 1] & ok)[0]
        if not len(fi) or not len(pi):
            continue
        Z = blk[:, np.union1d(fi, pi)]
        allidx = np.union1d(fi, pi)
        Z = np.nan_to_num((Z - np.nanmean(Z, 0)) / np.nanstd(Z, 0))
        pos = {c: j for j, c in enumerate(allidx)}
        A = Z[:, [pos[c] for c in fi]]
        B = Z[:, [pos[c] for c in pi]]
        C = A.T @ B / win
        C[fi[:, None] == pi[None, :]] = -np.inf
        top = np.argsort(-C, axis=1)[:, :k]
        out[me] = {int(f): pi[t] for f, t in zip(fi, top)}
    return out


def customer_links(F, focal: pd.DataFrame, files: list[str], valid_days: int = 456) -> dict:
    col = {t: i for i, t in enumerate(F["cols"])}
    d = []
    for f in files:
        if os.path.exists(f):
            x = pd.read_csv(f, parse_dates=["filed"])
            d.append(x[x["n_customer"] > 0][["filer", "target", "filed"]])
    d = pd.concat(d).drop_duplicates()
    d = d[d["filer"].isin(col) & d["target"].isin(col) & (d["filer"] != d["target"])]
    out = {}
    fo, ma = focal, F["mall"]
    for me in month_ends(F["dates"]):
        live = d[(d["filed"] <= me) & (d["filed"] > me - pd.Timedelta(days=valid_days))]
        g = {}
        for f, t in zip(live["filer"], live["target"]):
            if fo.at[me, f] and ma.at[me, t]:
                g.setdefault(col[f], set()).add(col[t])
        if g:
            out[me] = {i: np.array(sorted(s)) for i, s in g.items()}
    return out


def link_signal(F, links: dict, h: int, kind: str) -> pd.DataFrame:
    lr = np.log1p(F["ret"].fillna(0.0))
    R = lr.rolling(h, min_periods=h).sum().to_numpy()
    out = np.full(R.shape, np.nan)
    mes = sorted(links)
    for a, b in zip(mes, mes[1:] + [F["dates"][-1] + pd.Timedelta(days=1)]):
        rows = np.where((F["dates"] > a) & (F["dates"] <= b))[0]
        for i, nb in links[a].items():
            out[rows, i] = np.nanmean(R[np.ix_(rows, nb)], axis=1)
    S = pd.DataFrame(out, index=F["dates"], columns=F["cols"])
    if kind == "gap":
        S = S - pd.DataFrame(R, index=F["dates"], columns=F["cols"])
    return S


# ---------------------------------------------------------------------------
# measures
# ---------------------------------------------------------------------------

def _nw_t(x: pd.Series, lags: int) -> float:
    x = x.dropna()
    n = len(x)
    if n < 6:
        return np.nan
    u = x - x.mean()
    s = (u ** 2).sum() / n
    for l in range(1, lags + 1):
        s += 2 * (1 - l / (lags + 1)) * (u[l:].to_numpy() * u[:-l].to_numpy()).sum() / n
    return float(x.mean() / np.sqrt(s / n))


def forward(F, h: int) -> pd.DataFrame:
    o = F["open"].where(F["open"].notna(), F["close"]).ffill(limit=3)
    return o.shift(-(h + 1)) / o.shift(-1) - 1


def cross_section_tests(F, sig: pd.DataFrame, focal: pd.DataFrame, h: int, fwd: pd.DataFrame,
                        own21: pd.DataFrame, mom: pd.DataFrame, cost: float = 0.002) -> dict:
    """FM slope, quintile spread and top-quintile excess, sampled every 21 days."""
    rows = []
    idx = np.arange(260, len(F["dates"]) - h - 1, 21)
    for i in idx:
        d = F["dates"][i]
        f = focal.iloc[i]
        df = pd.DataFrame({"y": fwd.iloc[i], "s": sig.iloc[i], "own": own21.iloc[i], "m": mom.iloc[i]})[f.values]
        uni = df["y"].dropna()
        df = df.dropna()
        if len(df) < 60:
            continue
        rk = df.rank(pct=True) - 0.5
        X = np.column_stack([np.ones(len(rk)), rk[["s", "own", "m"]].to_numpy()])
        b, *_ = np.linalg.lstsq(X, rk["y"].to_numpy(), rcond=None)
        q = pd.qcut(df["s"].rank(method="first"), 5, labels=False)
        top, bot = df["y"][q == 4], df["y"][q == 0]
        rows.append({"date": d, "b": b[1], "ls": top.mean() - bot.mean(),
                     "top_gross": top.mean() - uni.mean(), "names": frozenset(top.index), "n": len(df)})
    T = pd.DataFrame(rows).set_index("date")
    out = {}
    lags = max(h // 21 - 1, 0)
    per_year = 252 / h
    for name, (a, z) in SPANS.items():
        x = T.loc[a:z]
        # non-overlapping portfolio periods for the h-day returns: every (h/21)-th sample
        step = max(h // 21, 1)
        xp = x.iloc[::step].copy()
        prev = [None] + list(xp["names"].iloc[:-1])
        turn = [1.0 if p_ is None else len(n_ - p_) / max(len(n_), 1) for n_, p_ in zip(xp["names"], prev)]
        xp["top_ex"] = xp["top_gross"] - 2 * cost * np.array(turn)
        out[name] = {"fm_b": float(x["b"].mean()), "fm_t": _nw_t(x["b"], lags),
                     "ls_ann": float(xp["ls"].mean() * per_year), "ls_t": _nw_t(xp["ls"], 0),
                     "top_ex_ann": float(xp["top_ex"].mean() * per_year), "top_ex_t": _nw_t(xp["top_ex"], 0),
                     "n": float(x["n"].mean()), "periods": len(x)}
    return out


def sleeve(F, sig: pd.DataFrame, h: int) -> dict:
    from research.lab.alpha import alpha_stats, bench_returns
    from research.lab.us_sweep import run
    ps = F["ps"]
    score = sig.reindex(index=ps.dates, columns=ps.close.columns).where(ps.member)
    r = run(ps, score, [("2017-01-03", END)], top_n=30, buffer=2.0, every=h)["_rets"]
    bq = bench_returns(r.index)
    bm = ps.benchmark.reindex(r.index).pct_change(fill_method=None) if ps.benchmark is not None else None
    out = {}
    for name, (a, z) in SPANS.items():
        x = r.loc[a:z]
        aq = alpha_stats(x, bq)
        am = alpha_stats(x, bm) if bm is not None else {"alpha": np.nan, "t": np.nan}
        out[name] = {"alpha_mid": am["alpha"], "t_mid": am["t"], "alpha_qqq": aq["alpha"], "t_qqq": aq["t"],
                     "cagr": aq["cagr"], "bench_mid_cagr": am.get("bench_cagr", np.nan)}
    return out, r


def events(F, big_links: dict, e_abn: np.ndarray) -> dict:
    ev = pd.read_csv(os.path.join(DIR, "sec_8k_earnings.csv"))
    ev["t"] = pd.to_datetime(ev["acceptanceDateTime"], utc=True).dt.tz_convert("America/New_York")
    ev["d"] = ev["t"].dt.normalize().dt.tz_localize(None) + pd.to_timedelta((ev["t"].dt.hour >= 16).astype(int), "D")
    dates = F["dates"]
    col = {t: i for i, t in enumerate(F["cols"])}
    pos = dates.searchsorted(ev["d"])
    ev = ev[(pos < len(dates))].assign(i=pos[pos < len(dates)])
    ev = ev[ev["ticker"].isin(col) & (ev["d"] >= pd.Timestamp("2016-06-01"))]
    mes = np.array(sorted(big_links))
    # invert: big j -> small followers, per month end
    inv = {me: {} for me in mes}
    for me in mes:
        for i, nb in big_links[me].items():
            for j in nb:
                inv[me].setdefault(int(j), []).append(i)
    mb = F["mb"].to_numpy()
    rows = []
    for t, i in zip(ev["ticker"], ev["i"]):
        if i + 64 >= len(dates):
            continue
        j = col[t]
        if not mb[i, j] or not np.isfinite(e_abn[i, j]):
            continue
        me = mes[mes < dates[i]]
        if not len(me):
            continue
        for q in inv[me[-1]].get(j, []):
            x = e_abn[i:i + 64, q]
            if not np.isfinite(x[0]):
                continue
            x = np.nan_to_num(x)
            rows.append((dates[i], t, F["cols"][q], e_abn[i, j], x[0], x[1:22].sum(), x[1:64].sum()))
    E = pd.DataFrame(rows, columns=["date", "announcer", "small", "ar_j", "ar0", "car1_21", "car1_63"])
    E.to_pickle(os.path.join(RESULTS, "lead_lag_small_events.pkl"))
    out = {}
    for name, (a, z) in SPANS.items():
        sub = E[(E["date"] >= a) & (E["date"] <= z)]
        res = {"pairs": len(sub)}
        for y in ("ar0", "car1_21", "car1_63"):
            g = sub.groupby(sub["date"].dt.to_period("M"))
            sl = g.apply(lambda d: np.polyfit(d["ar_j"], d[y], 1)[0] if len(d) > 20 and d["ar_j"].std() > 0 else np.nan).dropna()
            lags = 0 if y != "car1_63" else 2
            res[y] = {"slope": float(sl.mean()), "t": _nw_t(sl, lags)}
        out[name] = res
    return out


def main() -> None:
    os.makedirs(RESULTS, exist_ok=True)
    F = frame()
    e = residuals(F)
    c = F["close"].ffill()
    own21 = c / c.shift(21) - 1
    mom = c.shift(21) / c.shift(252) - 1
    fwd = {h: forward(F, h) for h in HS}
    ms, mb, mall = F["ms"], F["mb"], F["mall"]
    cov = pd.DataFrame({"small_members": ms.sum(axis=1), "big_members": mb.sum(axis=1)}).resample("YE").mean()
    print("members per day (yearly mean):\n", cov.loc["2016":].round(0).to_string(), flush=True)

    L = {"peer": peer_links(F, e, ms, mall, 10), "bigpeer": peer_links(F, e, ms, mb, 5),
         "peer_sp500": peer_links(F, e, mb, mall, 10)}
    have_mid = os.path.exists(os.path.join(DIR, "mid_10k_links.csv"))
    if have_mid:
        L["cust"] = customer_links(F, ms, [os.path.join(DIR, "mid_10k_links.csv"), os.path.join(DIR, "sec_10k_links.csv")])
    L["cust_sp500"] = customer_links(F, mb, [os.path.join(DIR, "sec_10k_links.csv")])
    for k, v in L.items():
        n = np.mean([len(x) for x in list(v.values()) if True][-90:]) if v else 0
        print(f"links {k}: {len(v)} months, {n:.0f} focal stocks with links (last 90 months avg)", flush=True)

    variants = []
    for h in HS:
        variants += [("peer_mom", "peer", "mom", h, "small"), ("peer_gap", "peer", "gap", h, "small"),
                     ("bigpeer_mom", "bigpeer", "mom", h, "small")]
        if have_mid:
            variants += [("cust_mom", "cust", "mom", h, "small"), ("cust_gap", "cust", "gap", h, "small")]
        variants += [("sp500_peer_mom", "peer_sp500", "mom", h, "big"), ("sp500_peer_gap", "peer_sp500", "gap", h, "big"),
                     ("sp500_cust_mom", "cust_sp500", "mom", h, "big")]
    res, rets, trials = {}, {}, []
    for name, lk, kind, h, uni in variants:
        sig = link_signal(F, L[lk], h, kind)
        focal = ms if uni == "small" else mb
        cs = cross_section_tests(F, sig, focal, h, fwd[h], own21, mom)
        row = {"cs": cs}
        if uni == "small":
            sl, r = sleeve(F, sig, h)
            row["sleeve"] = sl
            rets[f"{name}_{h}"] = r
        res[f"{name}_{h}"] = row
        m = cs["2019-26"]
        line = (f"{name:16s} h{h:<3d} FM t 19-26 {m['fm_t']:5.2f} | 19-22 {cs['2019-22']['fm_t']:5.2f} | 23-26 "
                f"{cs['2023-26']['fm_t']:5.2f} | 17-26 {cs['2017-26']['fm_t']:5.2f} | L-S {m['ls_ann']*100:5.1f}%/yr "
                f"(t {m['ls_t']:.2f}) | top-EW net {m['top_ex_ann']*100:5.1f}%/yr (t {m['top_ex_t']:.2f}) n {m['n']:.0f}")
        if "sleeve" in row:
            s = row["sleeve"]["2019-26"]
            line += f" | top30 a_mid {s['alpha_mid']*100:5.1f}% (t {s['t_mid']:.2f}) a_qqq {s['alpha_qqq']*100:5.1f}% (t {s['t_qqq']:.2f})"
        print(line, flush=True)
        trials.append({"name": name, "h": h, "universe": uni, **{f"fm_t_{k}": v["fm_t"] for k, v in cs.items()},
                       "ls_ann_1926": m["ls_ann"], "top_ex_ann_1926": m["top_ex_ann"]})
    # event test: S&P 500 earnings -> small followers
    abn = (F["ret"].sub(F["mkt"], axis=0)).to_numpy()
    res["events"] = events(F, L["bigpeer"], abn)
    for k, v in res["events"].items():
        print(f"events {k}: pairs {v['pairs']} | " + " | ".join(f"{y} slope {v[y]['slope']:+.3f} t {v[y]['t']:.2f}"
                                                              for y in ("ar0", "car1_21", "car1_63")), flush=True)
    trials += [{"name": f"event_{y}", "h": 0, "universe": "small"} for y in ("ar0", "car1_21", "car1_63")]
    json.dump(res, open(os.path.join(RESULTS, "lead_lag_small.json"), "w"), indent=1, default=float)
    pd.DataFrame(rets).to_pickle(os.path.join(RESULTS, "lead_lag_small_returns.pkl"))
    pd.DataFrame(trials).to_csv(os.path.join(RESULTS, "lead_lag_small_trials.csv"), index=False)
    print(f"{len(trials)} trials", flush=True)


if __name__ == "__main__":
    main()