quant-trading/research/lab/rotation.py

"""
Walk-forward factor weighting: does choosing factors from their recent record
beat holding a fixed composite?

Each scheme sets factor weights at every month end from data available then
(trailing factor returns / ICs); the stock score is the weighted average of
the factors' cross-sectional percentile ranks. The resulting scores run
through the same engine (PIT members, next-open fills, phase averaging).

    PYTHONPATH=. uv run python -m research.lab.rotation --market us
"""

from __future__ import annotations

import argparse
import os
import pickle
import time
import warnings

import numpy as np
import pandas as pd

from research.lab.factor_scan import RESULTS, load, phase_avg
from research.lab.factors import FACTORS
from research.lab.report import forward_returns

warnings.filterwarnings("ignore")


# ---------------------------------------------------------------------------
# inputs
# ---------------------------------------------------------------------------

def factor_ranks(p, market: str, names: list[str] | None = None) -> dict[str, pd.DataFrame]:
    """Percentile ranks among PIT members, cached per market (float32)."""
    path = os.path.join("data", "research", f"ranks_{market}.pkl")
    cache = {}
    if os.path.exists(path):
        with open(path, "rb") as fh:
            cache = pickle.load(fh)
    mk = "us" if market == "us" else "cn"
    names = names or [n for n, f in FACTORS.items() if mk in f.markets]
    dirty = False
    for n in names:
        if n in cache:
            continue
        t = time.time()
        try:
            f = FACTORS[n](p).reindex(index=p.dates, columns=p.close.columns)
        except Exception as exc:  # noqa: BLE001
            print(f"skip {n}: {exc}", flush=True)
            continue
        cache[n] = f.where(p.member).rank(axis=1, pct=True).astype("float32")
        dirty = True
        print(f"  rank {n} [{time.time() - t:.0f}s]", flush=True)
    if dirty:
        with open(path, "wb") as fh:
            pickle.dump(cache, fh, protocol=4)
    return {n: cache[n] for n in names if n in cache}


def factor_returns(p, ranks: dict[str, pd.DataFrame], q: float = 0.2) -> pd.DataFrame:
    """Daily return of each factor's top quintile minus the member average,
    formed on the previous close's ranks (a timing input, not a tradable P&L)."""
    r = p.close.pct_change(fill_method=None)
    mem = p.member.shift(1, fill_value=False)
    base = r.where(mem).mean(axis=1)
    out = {}
    for n, rk in ranks.items():
        top = rk.shift(1) >= 1 - q
        out[n] = r.where(top & mem).mean(axis=1) - base
    return pd.DataFrame(out).fillna(0.0)


def monthly_ic(p, ranks: dict[str, pd.DataFrame], horizon: int = 21) -> pd.DataFrame:
    """Rank IC of each factor at month ends vs the next ``horizon`` days' open-to-open
    return. The IC of month m is only known ``horizon`` days later."""
    fwd = forward_returns(p, horizon).where(p.member).rank(axis=1, pct=True)
    me = _month_ends(p.dates)
    out = {}
    for n, rk in ranks.items():
        a = rk.loc[me]
        b = fwd.loc[me]
        am = a.sub(a.mean(axis=1), axis=0)
        bm = b.sub(b.mean(axis=1), axis=0)
        ok = a.notna() & b.notna()
        cov = (am * bm).where(ok).sum(axis=1)
        out[n] = cov / np.sqrt((am ** 2).where(ok).sum(axis=1) * (bm ** 2).where(ok).sum(axis=1))
    return pd.DataFrame(out)


def _month_ends(dates: pd.DatetimeIndex) -> pd.DatetimeIndex:
    s = pd.Series(dates, index=dates)
    return pd.DatetimeIndex(s.groupby(dates.to_period("M")).max().values)


# ---------------------------------------------------------------------------
# weighting schemes: each returns month-end x factor weights (>= 0)
# ---------------------------------------------------------------------------

def _monthly(fr: pd.DataFrame) -> pd.DataFrame:
    return (1 + fr).resample("ME").prod() - 1


def w_equal(fr: pd.DataFrame, names: list[str]) -> pd.DataFrame:
    return pd.DataFrame(1.0, index=_monthly(fr).index, columns=names)


def w_tsmom(fr: pd.DataFrame, names: list[str], lookback: int = 12) -> pd.DataFrame:
    """Time-series factor momentum (Gupta & Kelly 2019; Ehsani & Linnainmaa 2022):
    hold a factor when its trailing ``lookback``-month return is positive."""
    m = _monthly(fr[names])
    trail = (1 + m).rolling(lookback, min_periods=lookback).apply(np.prod, raw=True) - 1
    return (trail > 0).astype(float).where(trail.notna())


def w_csmom(fr: pd.DataFrame, names: list[str], lookback: int = 12, k: int = 5) -> pd.DataFrame:
    """Cross-sectional factor momentum: the ``k`` factors with the best trailing return."""
    m = _monthly(fr[names])
    trail = (1 + m).rolling(lookback, min_periods=lookback).apply(np.prod, raw=True) - 1
    rank = trail.rank(axis=1, ascending=False)
    return (rank <= k).astype(float).where(trail.notna())


def w_ic(ic: pd.DataFrame, names: list[str], lookback: int = 36) -> pd.DataFrame:
    """Weight by the trailing mean IC, clipped at zero. The IC of month end m
    needs 22 more sessions of prices, so it is used from two month ends later."""
    known = ic[names].shift(2)
    mean = known.rolling(lookback, min_periods=lookback // 2).mean()
    return mean.clip(lower=0)


def w_icir(ic: pd.DataFrame, names: list[str], lookback: int = 36) -> pd.DataFrame:
    known = ic[names].shift(2)
    roll = known.rolling(lookback, min_periods=lookback // 2)
    return (roll.mean() / roll.std()).clip(lower=0)


def w_state(fr: pd.DataFrame, names: list[str], state: pd.Series, min_obs: int = 24) -> pd.DataFrame:
    """Regime-conditional factor choice: in each month, weight factors whose
    average return in past months of the SAME market state was positive.
    ``state`` is a month-end label known at that month end."""
    m = _monthly(fr[names])
    st = state.reindex(m.index)
    out = pd.DataFrame(np.nan, index=m.index, columns=names)
    for i in range(len(m.index)):
        # weights set at month end i are used during month i+1; the history
        # pairs each month's return with the state at the end of the month before
        hist = m.iloc[:i + 1]
        prev_state = st.iloc[:i + 1].shift(1)
        same = hist[prev_state == st.iloc[i]]
        if len(same) >= min_obs:
            out.iloc[i] = (same.mean() > 0).astype(float).values
    return out


def market_state(p) -> pd.Series:
    """Month-end state: EW member index above/below its 10-month average x
    trailing 1-month vol above/below its 5-year median."""
    r = p.close.pct_change(fill_method=None).where(p.member.shift(1, fill_value=False)).mean(axis=1).fillna(0)
    lvl = (1 + r).cumprod()
    mlvl = lvl.resample("ME").last()
    trend = mlvl > mlvl.rolling(10, min_periods=10).mean()
    vol = r.rolling(21).std().resample("ME").last()
    hi = vol > vol.rolling(60, min_periods=24).median()
    lab = trend.map({True: "up", False: "down"}) + "_" + hi.map({True: "hv", False: "lv"})
    return lab.where(mlvl.rolling(10).count() >= 10)


# ---------------------------------------------------------------------------
# blend
# ---------------------------------------------------------------------------

def blend(p, ranks: dict[str, pd.DataFrame], W: pd.DataFrame) -> pd.DataFrame:
    """Daily score = weighted mean of available factor ranks, using the most
    recent month-end weights (applied from the next trading day)."""
    names = [c for c in W.columns if c in ranks]
    Wd = W[names].reindex(p.dates, method="ffill").shift(1)
    num = np.zeros(p.close.shape, dtype="float64")
    den = np.zeros(p.close.shape, dtype="float64")
    for n in names:
        R = ranks[n].values
        w = np.nan_to_num(Wd[n].values, nan=0.0)[:, None]
        ok = ~np.isnan(R)
        num += np.where(ok, R, 0.0) * w
        den += ok * w
    with np.errstate(invalid="ignore", divide="ignore"):
        s = num / den
    s[den < 1e-9] = np.nan
    return pd.DataFrame(s, index=p.dates, columns=p.close.columns).where(p.member)


# ---------------------------------------------------------------------------

EVAL = {
    "us": [("2010-01-01", "2013-06-30"), ("2013-07-01", "2016-09-26"),
           ("2016-09-27", "2021-09-26"), ("2021-09-27", "2026-09-24")],
    "cn": [("2012-01-01", "2014-06-30"), ("2014-07-01", "2016-09-26"),
           ("2016-09-27", "2021-09-26"), ("2021-09-27", "2026-09-24")],
}


def evaluate(p, score, subs, **kw) -> dict:
    per = [phase_avg(p, score, a, b, **kw) for a, b in subs]
    full = phase_avg(p, score, subs[0][0], subs[-1][1], **kw)
    return {"sharpe": full["sharpe"], "cagr": full["cagr"], "mdd": full["mdd"],
            "turn": full["turnover"], "min_block": min(x["sharpe"] for x in per),
            "blocks": " ".join(f"{x['sharpe']:.2f}" for x in per)}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--market", default="us")
    ap.add_argument("--top-n", type=int, default=20)
    ap.add_argument("--every", type=int, default=10)
    ap.add_argument("--buffer", type=float, default=1.0)
    args = ap.parse_args()
    mk = "us" if args.market == "us" else "cn"
    p, _ = load(args.market)
    t = time.time()
    ranks = factor_ranks(p, args.market)
    names = list(ranks)
    print(f"{len(names)} factors ranked [{time.time() - t:.0f}s]", flush=True)
    fr = factor_returns(p, ranks)
    ic = monthly_ic(p, ranks)
    with open(os.path.join(RESULTS, f"factor_returns_{args.market}.pkl"), "wb") as fh:
        pickle.dump({"fr": fr, "ic": ic}, fh)
    state = market_state(p)
    schemes = {
        "equal_all": w_equal(fr, names),
        "tsmom_12": w_tsmom(fr, names, 12),
        "tsmom_6": w_tsmom(fr, names, 6),
        "tsmom_1": w_tsmom(fr, names, 1),
        "csmom_12_top5": w_csmom(fr, names, 12, 5),
        "csmom_12_top10": w_csmom(fr, names, 12, 10),
        "csmom_1_top5": w_csmom(fr, names, 1, 5),
        "csmom_36_top5": w_csmom(fr, names, 36, 5),
        "ic_36": w_ic(ic, names, 36),
        "ic_60": w_ic(ic, names, 60),
        "icir_36": w_icir(ic, names, 36),
        "state_cond": w_state(fr, names, state),
    }
    subs = EVAL[mk]
    kw = dict(top_n=args.top_n, every=args.every, buffer=args.buffer)
    rows = []
    for name, W in schemes.items():
        s = blend(p, ranks, W)
        st = evaluate(p, s, subs, **kw)
        rows.append({"scheme": name, **st})
        nact = (W.loc[subs[0][0]:] > 0).sum(axis=1).mean()
        print(f"{name:16s} sh {st['sharpe']:.2f} cagr {st['cagr']*100:5.1f}% mdd {st['mdd']*100:5.1f}% "
              f"turn {st['turn']:4.1f}x min {st['min_block']:.2f} [{st['blocks']}] active {nact:.1f}", flush=True)
    pd.DataFrame(rows).to_csv(os.path.join(RESULTS, f"rotation_{args.market}.csv"), index=False)


if __name__ == "__main__":
    main()