quant-trading/research/lab/theme_propagation.py

"""
Does a main uptrend propagate along linked industries with a lag of months?
The AI build-out moved stage by stage (hyperscalers and power in H1 2025,
optics / equipment / memory in H2 2025, chips and servers in Q2 2026). This
asks whether that pattern is systematic: when an industry is in a strong
run, do the industries linked to it that have NOT moved yet outperform over
the next 1-6 months? Every era gets its own theme; the rule never names one.

    PYTHONPATH=. uv run python -m research.lab.theme_propagation

Units: 3-digit SIC industries of PIT S&P 500 members (sec_sic.csv), equal
weight, with at least 3 members that month. Monthly, at each month end t:

  hot       industries whose 6-month return over the equal-weight market is
            in the top 10% of industries and above +15%
  links     econ_up   supplier industries of the hot one: 10-K customer
                      sentences filed in the last 3 years (filer = supplier,
                      named company = customer; research/lab/fetch_supply_links.py)
            econ_any  any 10-K mention in either direction, last 3 years
            corr      the 5 industries whose monthly excess returns correlated
                      most with it over the last 36 months
            lagged    the 5 industries whose excess return in months m+1..m+3
                      correlated most with its return in month m, last 60 months
  followers linked industries that are not hot themselves and whose own
            6-month excess return is below the cross-industry median
            ("the linked industry has not moved yet"); variant: all linked
  control   below-median industries not linked to any hot industry

Portfolio: equal weight over the member stocks of the follower industries,
bought at the first open after t and held 1, 3 or 6 months (overlapping
cohorts averaged, i.e. monthly tranches), against the equal-weight market.
20 bp per unit of turnover is charged at each rebalance. Reported per era
(2006-10, 2011-15, 2016-20, 2021-26) and design (<=2016) / holdout.
A stock-level check follows: the lead_lag link signals over 63- and 126-day
windows predicting the next 63 days (Fama-MacBeth with the usual controls).
"""

from __future__ import annotations

import json
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.factor_scan import RESULTS, load

warnings.filterwarnings("ignore")
ERAS = [("2006", "2010"), ("2011", "2015"), ("2016", "2020"), ("2021", "2026")]
HOT_Q, HOT_MIN, LOOK = 0.90, 0.15, 6


def industry_map(p) -> pd.Series:
    sic = pd.read_csv(os.path.join("data", "research", "sec_sic.csv")).drop_duplicates("ticker").set_index("ticker")["sic"]
    s = (pd.to_numeric(sic, errors="coerce") // 10).dropna().astype(int)
    return s.reindex(p.close.columns)


def monthly_open_returns(p) -> tuple[pd.DataFrame, pd.DatetimeIndex]:
    """Stock return from the first open of month m+1 to the first open of
    month m+2, indexed by the signal month end m."""
    o = p.open.where(p.open.notna(), p.close).ffill()
    first = pd.Series(p.dates, index=p.dates).groupby(p.dates.to_period("M")).min()
    ends = pd.Series(p.dates, index=p.dates).groupby(p.dates.to_period("M")).max()
    fo = o.loc[first.values]
    r = fo.shift(-1) / fo - 1                       # first open of month k -> first open of month k+1
    r.index = ends.values
    r = r.shift(-1)                                 # signal at end of m -> return over month m+1
    return r, pd.DatetimeIndex(ends.values)


def group_panels(p, ind: pd.Series):
    """Monthly member masks and equal-weight monthly close-to-close excess
    returns per industry (for the signals)."""
    ends = pd.Series(p.dates, index=p.dates).groupby(p.dates.to_period("M")).max().values
    c = p.close.ffill()
    mc = c.loc[ends]
    mr = (mc / mc.shift(1) - 1).where(p.member.loc[ends].shift(1).fillna(False).astype(bool))
    mem = p.member.loc[ends]
    mkt = mr.mean(axis=1)
    groups = sorted(ind.dropna().unique())
    G = pd.DataFrame(index=mr.index, columns=groups, dtype=float)
    N = pd.DataFrame(index=mr.index, columns=groups, dtype=float)
    for g in groups:
        cols = list(ind.index[ind == g])
        G[g] = mr[cols].mean(axis=1)
        N[g] = mem[cols].sum(axis=1)
    X = G.sub(mkt, axis=0)
    return X.where(N >= 3), N, mem


def econ_links(ind: pd.Series, ends, role: str) -> dict:
    """{month end: set of (hot industry, linked industry)} from 10-Ks filed
    in the last 3 years. role "up": customer industry -> supplier industry."""
    d = pd.read_csv(os.path.join("data", "research", "sec_10k_links.csv"), parse_dates=["filed"])
    if role == "up":
        d = d[d["n_customer"] > 0]
    d = d.assign(fi=d["filer"].map(ind), ti=d["target"].map(ind)).dropna(subset=["fi", "ti"])
    d = d[d["fi"] != d["ti"]]
    out = {}
    for t in ends:
        live = d[(d["filed"] <= t) & (d["filed"] > t - pd.Timedelta(days=3 * 365))]
        if role == "up":            # filer supplies the named customer: customer industry leads
            pairs = set(zip(live["ti"].astype(int), live["fi"].astype(int)))
        else:
            pairs = set(zip(live["ti"].astype(int), live["fi"].astype(int))) | set(zip(live["fi"].astype(int), live["ti"].astype(int)))
        out[t] = pairs
    return out


def stat_links(X: pd.DataFrame, kind: str, k: int = 5) -> dict:
    out = {}
    for i, t in enumerate(X.index):
        win = 36 if kind == "corr" else 60
        if i < win:
            out[t] = set()
            continue
        W = X.iloc[i - win + 1:i + 1]
        W = W.loc[:, W.notna().sum() >= int(win * 0.8)]
        if W.shape[1] < 10:
            out[t] = set()
            continue
        Z = (W - W.mean()) / W.std()
        Z = Z.fillna(0.0)
        if kind == "corr":
            C = Z.T @ Z / len(Z)
        else:                       # C[a, b] = mean over lags 1-3 of corr(a in month m, b in month m+lag)
            C = sum(Z.iloc[:-lag].T.to_numpy() @ Z.iloc[lag:].to_numpy() / (len(Z) - lag) for lag in (1, 2, 3)) / 3
            C = pd.DataFrame(C, index=Z.columns, columns=Z.columns)
        C = C.to_numpy(copy=True)
        np.fill_diagonal(C, -np.inf)
        cols = list(Z.columns)
        pairs = set()
        for a in range(len(cols)):
            for b in np.argsort(-C[a])[:k]:
                pairs.add((cols[a], cols[b]))
        out[t] = pairs
    return out


def followers(X: pd.DataFrame, links: dict, laggard_only: bool = True) -> tuple[dict, dict, dict]:
    """Per month end: hot industries, follower industries, control industries."""
    cum = np.log1p(X.fillna(0.0)).rolling(LOOK, min_periods=LOOK).sum().where(X.notna())
    hot_d, fol_d, ctl_d = {}, {}, {}
    for t in X.index:
        row = cum.loc[t].dropna()
        if len(row) < 20:
            continue
        thr = max(row.quantile(HOT_Q), np.log1p(HOT_MIN))
        hot = set(row.index[row >= thr])
        med = row.median()
        lag = set(row.index[row < med])
        pairs = links.get(t, set())
        linked = {b for a, b in pairs if a in hot} - hot
        linked &= set(row.index)
        fol = linked & lag if laggard_only else linked
        hot_d[t], fol_d[t] = hot, fol
        ctl_d[t] = lag - linked - hot
    return hot_d, fol_d, ctl_d


def portfolio(sel: dict, ind: pd.Series, mem: pd.DataFrame, R: pd.DataFrame, hold: int) -> pd.Series:
    """Monthly return of equal-weight member stocks in the selected industries,
    held ``hold`` months as overlapping cohorts, net of 20 bp x turnover."""
    months = R.index
    W = pd.DataFrame(0.0, index=months, columns=R.columns)
    for t, gs in sel.items():
        if t not in W.index or not gs:
            continue
        cols = [c for c in ind.index[ind.isin(gs)] if mem.at[t, c]]
        if cols:
            W.loc[t, cols] = 1.0 / len(cols)
    Wh = sum(W.shift(k).fillna(0.0) for k in range(hold)) / hold     # tranches formed at t, t-1, ...
    invested = Wh.sum(axis=1)
    gross = (Wh * R.fillna(0.0)).sum(axis=1) / invested.replace(0, np.nan)
    turn = (Wh.div(invested.replace(0, np.nan), axis=0).fillna(0.0).diff().abs().sum(axis=1)) / 2
    return (gross - 0.002 * 2 * turn).where(invested > 0)


def mstats(r: pd.Series, b: pd.Series) -> dict:
    """Excess over the benchmark: mean per month, annualized, Newey-West t (3 lags)."""
    x = (r - b).dropna()
    if len(x) < 12:
        return {"ann": np.nan, "t": np.nan, "n": len(x)}
    n, m = len(x), x.mean()
    e = x - m
    s = (e ** 2).sum() / n
    for l in (1, 2, 3):
        s += 2 * (1 - l / 4) * (e.iloc[l:].to_numpy() * e.iloc[:-l].to_numpy()).sum() / n
    return {"ann": float(m * 12), "t": float(m / np.sqrt(s / n)), "n": n}


def stock_level_long(p) -> dict:
    """lead_lag link signals over long windows, next-63-day Fama-MacBeth."""
    from research.lab import lead_lag as ll
    from research.lab.rotation import factor_ranks
    pr = factor_ranks(p, "us", ["overnight_mom", "mom_12_1"])
    L = pickle.load(open(ll.CACHE, "rb"))
    c = p.close.ffill()
    out = {}
    for win in (63, 126):
        own = (c / c.shift(win) - 1).where(p.member)
        for key in ("corr_peer", "10k_any", "10k_customer"):
            for kind in ("mom", "gap"):
                sig = ll.link_signal(p, L[key], win, kind)
                fm = ll.fama_macbeth(p, sig, {"own": own, "mom_12_1": pr["mom_12_1"], "overnight": pr["overnight_mom"]}, 63)
                out[f"{key}_{kind}_{win}d"] = fm
                print(f"  {key}_{kind} window {win}d -> next 63d: FM t all {fm['all']['t']:5.2f} | design {fm['design']['t']:5.2f} "
                      f"| holdout {fm['holdout']['t']:5.2f} | IC {fm['all']['ic']:+.3f}", flush=True)
    return out


def main() -> None:
    p, _ = load("us")
    ind = industry_map(p)
    X, N, mem = group_panels(p, ind)
    R, ends = monthly_open_returns(p)            # unmasked: a stock bought at t keeps its return after leaving the index
    mkt = R.where(mem).mean(axis=1)
    from research.lab.alpha import bench_returns
    q = bench_returns(p.dates)
    qm = (1 + q.fillna(0.0)).groupby(p.dates.to_period("M")).prod() - 1
    qm.index = ends
    qm = qm.shift(-1)                      # align with R (return over month m+1)
    print(f"{X.notna().any().sum()} industries with >=3 members at some point; median per month "
          f"{int(X.notna().sum(axis=1).median())}", flush=True)
    link_sets = {"econ_up": econ_links(ind, X.index, "up"), "econ_any": econ_links(ind, X.index, "any"),
                 "corr": stat_links(X, "corr"), "lagged": stat_links(X, "lagged")}
    res, rets = {}, {}
    for name, L in link_sets.items():
        for lag_only in (True, False):
            hot, fol, ctl = followers(X, L, lag_only)
            nf = np.mean([len(v) for v in fol.values()])
            tag = f"{name}{'' if lag_only else ' (all linked)'}"
            for hold in (1, 3, 6):
                r = portfolio(fol, ind, mem, R, hold)
                rc = portfolio(ctl, ind, mem, R, hold)
                row = {"followers_per_month": float(nf)}
                for lab, a, b in [("all", "2006-09", "2026-08"), ("design", "2006-09", "2016-12"), ("holdout", "2017-01", "2026-08")] + \
                        [(f"{e0}-{e1}", e0, e1) for e0, e1 in ERAS]:
                    row[lab] = mstats(r.loc[a:b], mkt.loc[a:b])
                    row[lab + "_vs_control"] = mstats(r.loc[a:b], rc.loc[a:b])
                row["vs_qqq_holdout"] = mstats(r.loc["2017":"2026-08"], qm.loc["2017":"2026-08"])
                res[f"{tag} hold{hold}"] = row
                rets[f"{tag} hold{hold}"] = r
                eras = " ".join(f"{k}: {row[k]['ann']*100:+5.1f}% (t {row[k]['t']:+.2f})" for k in [f"{a}-{b}" for a, b in ERAS])
                print(f"{tag:22s} hold {hold}m  n {nf:4.1f} | vs mkt design {row['design']['ann']*100:+5.1f}% t {row['design']['t']:+.2f} "
                      f"holdout {row['holdout']['ann']*100:+5.1f}% t {row['holdout']['t']:+.2f} | vs control holdout "
                      f"{row['holdout_vs_control']['ann']*100:+5.1f}% t {row['holdout_vs_control']['t']:+.2f} | {eras}", flush=True)
    # which themes fired: most frequent hot industries per era
    hot, _, _ = followers(X, link_sets["corr"], True)
    sic = pd.read_csv(os.path.join("data", "research", "sec_sic.csv"))
    desc = sic.assign(g=pd.to_numeric(sic["sic"], errors="coerce") // 10).groupby("g")["sicDescription"].first()
    themes = {}
    for e0, e1 in ERAS:
        cnt = pd.Series([g for t, hs in hot.items() if e0 <= str(t.year) <= e1 for g in hs]).value_counts().head(6)
        themes[f"{e0}-{e1}"] = {f"{int(g)} {desc.get(g, '')}": int(n) for g, n in cnt.items()}
        print(f"hot industries {e0}-{e1}: " + ", ".join(f"{k} ({v}m)" for k, v in themes[f"{e0}-{e1}"].items()), flush=True)
    print("\nstock level, long windows:", flush=True)
    res["_stock_level"] = stock_level_long(p)
    res["_themes"] = themes
    json.dump(res, open(os.path.join(RESULTS, "theme_propagation.json"), "w"), indent=1, default=float)
    pd.DataFrame(rets).to_pickle(os.path.join(RESULTS, "theme_propagation_returns.pkl"))


if __name__ == "__main__":
    main()