quant-trading/research/lab/attention.py

"""
Retail attention from Wikipedia pageviews: does it predict S&P 500 returns,
and does it improve the momentum sleeve?

    uv run python -m research.lab.fetch_wiki_attention      # data/research/wiki_pageviews.pkl
    PYTHONPATH=. uv run --with numba python -m research.lab.attention [--views views_main]

Why Wikipedia: it is the only free, daily, point-in-time attention series that
covers every S&P 500 name back to 2015-07. X full-archive search needs
enterprise/pay-per-use access; Reddit's own API has no history (Pushshift is
moderator-only; Arctic Shift's full-text aggregates time out on
r/wallstreetbets; the academictorrents dumps are the remaining route);
StockTwits closed developer sign-up; Google Trends caps daily data at
~9-month windows with per-query scaling.

Timing: a pageview day is a UTC day, which ends at 20:00 ET, after the close.
The signal on trading day t uses views through UTC day t-1 only; trades fill
at the open of t+1 (research engine convention).

Signals (x = log(1 + daily views), article + its redirects):
  abn_att    mean x over the last 5 days minus the median x of the 252 days before
             (Da, Engelberg & Gao 2011 abnormal search volume, on Wikipedia)
  att_shock  x on the last day minus mean x over the last 21 days
  att_trend  mean x over the last 63 days minus mean x over the 189 days before
  att_level  mean x over the last 252 days (popularity; mostly a size proxy)

Tests, fixed before looking at results:
  (a) daily cross-sectional rank IC with the next-open to next-open 5- and
      21-day return; t from non-overlapping samples.
  (b) top-20 sleeves, each signal high and low (8 trials): abn_att/att_shock
      every 5 days, att_trend/att_level every 21 days, buffer 1. Alpha vs QQQ
      (Newey-West), design 2016-07..2020-12, holdout 2021..2026-09.
  (c) the momentum sleeve (overnight_mom + rec_126, top 10, every 10 days,
      buffer 2 = live overnight_recovery_top10) with 5 attention rules:
        veto_hi_abn     drop names in the top 10% of abn_att
        veto_hi_shock   drop names in the top 10% of att_shock
        resort_low_abn  top 30 by momentum, keep the 10 with the lowest abn_att
        resort_high_abn top 30 by momentum, keep the 10 with the highest abn_att
        resort_trend    top 30 by momentum, keep the 10 with the highest att_trend
      A name without pageviews is never vetoed and is neutral in a re-sort.
"""

from __future__ import annotations

import dataclasses
import json
import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.alpha import _nw_se, alpha_stats, bench_returns
from research.lab.factor_scan import RESULTS, load
from research.lab.report import stats
from research.lab.rotation import factor_ranks
from research.lab.us_sweep import _blend, log_trials, run

warnings.filterwarnings("ignore")
VIEWS = os.path.join("data", "research", "wiki_pageviews.pkl")
SPAN = ("2016-07-01", "2026-09-24")          # first year of views is warm-up
DESIGN = ("2016-07-01", "2020-12-31")
HOLDOUT = ("2021-01-01", "2026-09-24")


def load_views(p, which: str = "views") -> pd.DataFrame:
    """which = "views" (article + redirects) or "views_main" (article only)."""
    with open(VIEWS, "rb") as fh:
        d = pickle.load(fh)
    print(f"{which}; redirects fetched: {d.get('complete', True)}")
    return d[which].reindex(columns=p.close.columns)


def signals(p, views: pd.DataFrame) -> dict[str, pd.DataFrame]:
    """Calendar-day signals, lagged one day, sampled at each trading close."""
    x = np.log1p(views)
    base = x.shift(5).rolling(252, min_periods=200).median()
    sig = {
        "abn_att": x.rolling(5, min_periods=4).mean() - base,
        "att_shock": x - x.rolling(21, min_periods=15).mean(),
        "att_trend": x.rolling(63, min_periods=50).mean() - x.shift(63).rolling(189, min_periods=150).mean(),
        "att_level": x.rolling(252, min_periods=200).mean(),
    }
    out = {}
    for k, s in sig.items():
        s = s.shift(1)                                  # through UTC day t-1
        s = s.reindex(s.index.union(p.dates)).ffill(limit=3).reindex(p.dates)
        out[k] = s.where(p.member)
    return out


def forward(p, h: int) -> pd.DataFrame:
    o = p.open.where(p.open.notna(), p.close.ffill())
    return o.shift(-1 - h) / o.shift(-1) - 1


def ic_table(p, S: dict[str, pd.DataFrame]) -> pd.DataFrame:
    rows = []
    for h in (5, 21):
        f = forward(p, h).where(p.member)
        for k, s in S.items():
            s2, f2 = s.where(f.notna()), f.where(s.notna())
            ic = s2.rank(axis=1).corrwith(f2.rank(axis=1), axis=1, method="pearson")
            ic = ic.loc[SPAN[0]:SPAN[1]].dropna()
            for label, (a, z) in (("full", SPAN), ("design", DESIGN), ("holdout", HOLDOUT)):
                x = ic.loc[a:z]
                nov = x.iloc[::h]
                rows.append({"signal": k, "h": h, "period": label, "ic": x.mean(),
                             "t": nov.mean() / nov.std() * np.sqrt(len(nov)), "n": len(nov),
                             "hit": (x > 0).mean()})
    return pd.DataFrame(rows)


def paired_t(d: pd.Series) -> float:
    d = d.dropna()
    se = _nw_se(np.ones((len(d), 1)), (d - d.mean()).to_numpy())[0]
    return float(d.mean() / se)


def line(r: pd.Series, b: pd.Series) -> tuple[str, dict]:
    rec, cells = {}, []
    for label, (a, z) in (("full", SPAN), ("design", DESIGN), ("holdout", HOLDOUT)):
        x = r.loc[a:z]
        s = stats((1 + x).cumprod())
        al = alpha_stats(x, b)
        rec[label] = {"cagr": s["cagr"], "sharpe": s["sharpe"], "mdd": s["mdd"],
                      "alpha": al["alpha"], "t": al["t"], "beta": al["beta"]}
        cells.append(f"{s['cagr']*100:5.1f}% sh {s['sharpe']:.2f} a {al['alpha']*100:5.1f}% t {al['t']:5.2f}")
    return " | ".join(cells), rec


def main() -> None:
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--views", default="views", choices=["views", "views_main"])
    args = ap.parse_args()
    tag = "" if args.views == "views" else "_main"
    p, _ = load("us")
    views = load_views(p, args.views)
    b = bench_returns(p.dates)
    S = signals(p, views)
    cov = S["abn_att"].notna().sum(axis=1) / p.member.sum(axis=1)
    print("coverage of PIT members with abn_att:",
          {y: round(float(cov.loc[str(y)].mean()), 3) for y in range(2016, 2027)})

    out = {"coverage": {str(y): float(cov.loc[str(y)].mean()) for y in range(2016, 2027)}}
    ic = ic_table(p, S)
    print("\n(a) rank IC (next-open fwd return), t from non-overlapping samples")
    print(ic.pivot_table(index=["signal", "h"], columns="period", values=["ic", "t"]).round(4).to_string())
    out["ic"] = ic.to_dict("records")

    print(f"\n{'':26s}{'2016-07..2026-09':>40s} | {'design 16-20':>40s} | {'holdout 21-26':>40s}")
    R, rows = {}, []
    every = {"abn_att": 5, "att_shock": 5, "att_trend": 21, "att_level": 21}
    for k, s in S.items():
        for side, sc in (("high", s), ("low", -s)):
            st = run(p, sc.where(p.member), [SPAN], top_n=20, buffer=1.0, every=every[k])
            name = f"{k} {side} top20"
            R[name] = st["_rets"]
            rows.append({"score": f"attention:{name}", **{kk: v for kk, v in st.items() if not kk.startswith("_")}})
            txt, rec = line(st["_rets"], b)
            out[name] = rec
            print(f"{name:26s}{txt}", flush=True)
    ew = p.close.pct_change(fill_method=None).where(p.member.shift(1, fill_value=False)).mean(axis=1)
    txt, out["EW S&P 500"] = line(ew, b)
    print(f"{'EW S&P 500 (reference)':26s}{txt}")

    # (c) momentum interaction
    pr = factor_ranks(p, "us", ["overnight_mom", "rec_126"])
    mom = _blend(pr["overnight_mom"], pr["rec_126"]).where(p.member)
    m = p.member
    pct = {k: S[k].rank(axis=1, pct=True) for k in ("abn_att", "att_shock", "att_trend")}
    mom_pct = mom.rank(axis=1, pct=True)
    rc = {k: float(v.loc[SPAN[0]:].iloc[::21].corrwith(mom_pct.loc[SPAN[0]:].iloc[::21], axis=1).mean())
          for k, v in pct.items()}
    print("\nmean cross-sectional rank corr with the momentum score:", {k: round(v, 3) for k, v in rc.items()})
    out["corr_with_momentum"] = rc
    top30 = mom.rank(axis=1, ascending=False) <= 30
    ok = lambda veto: m & ~veto.fillna(False)                   # noqa: E731
    V = {
        "base": (mom, m, 2.0),
        "veto_hi_abn": (mom, ok(pct["abn_att"] > 0.9), 2.0),
        "veto_hi_shock": (mom, ok(pct["att_shock"] > 0.9), 2.0),
        "resort_low_abn": ((1 - pct["abn_att"]).fillna(0.5).where(top30), m, 1.0),
        "resort_high_abn": (pct["abn_att"].fillna(0.5).where(top30), m, 1.0),
        "resort_trend": (pct["att_trend"].fillna(0.5).where(top30), m, 1.0),
    }
    print(f"\n(c) momentum top10 with attention rules\n{'':26s}{'2016-07..2026-09':>40s} | {'design 16-20':>40s} | {'holdout 21-26':>40s} | vs base")
    M = {}
    for name, (score, elig, buf) in V.items():
        pv = dataclasses.replace(p, member=elig)
        st = run(pv, score.where(elig), [SPAN], top_n=10, buffer=buf, every=10)
        M[name] = st["_rets"]
        if name != "base":
            rows.append({"score": f"attention_mom:{name}", **{kk: v for kk, v in st.items() if not kk.startswith("_")}})
        txt, rec = line(st["_rets"], b)
        extra = ""
        if name != "base":
            d, dd = (M[name] - M["base"]).loc[HOLDOUT[0]:], (M[name] - M["base"]).loc[DESIGN[0]:DESIGN[1]]
            extra = (f"holdout {d.mean()*252*100:+5.1f}%/yr t {paired_t(d):+.2f} "
                     f"(design {dd.mean()*252*100:+5.1f}% t {paired_t(dd):+.2f})")
            rec["diff"] = {"holdout": float(d.mean() * 252), "holdout_t": paired_t(d),
                           "design": float(dd.mean() * 252), "design_t": paired_t(dd)}
        out["mom:" + name] = rec
        print(f"{name:26s}{txt} | {extra}", flush=True)
    if not tag:                                # the article-only run is a robustness check
        log_trials(rows, "attention")
        print(f"\ntrials logged: {len(rows)}")

    with open(os.path.join(RESULTS, f"attention{tag}.json"), "w") as fh:
        json.dump(out, fh, indent=1, default=float)
    pd.DataFrame({**R, **{"mom:" + k: v for k, v in M.items()}}).to_pickle(os.path.join(RESULTS, f"attention{tag}.pkl"))


if __name__ == "__main__":
    main()