scripts/backtest_c24.py

#!/usr/bin/env python3
"""C24: Add DJIA + VIX as macro features.

New features:
  djia_ret_1d     — prior-day DJIA return (non-leaking; US closes before TW opens)
  djia_ret_5d     — DJIA 5-day return
  djia_ma20_ratio — DJIA / 20d MA (trend regime)
  vix_level       — VIX spot level (fear index)
  vix_change_5d   — 5-day change in VIX

Hypothesis: DJIA captures broad US equity risk appetite beyond SOX/TNX;
VIX encodes the risk-off/risk-on regime that dominates TW next-day opens.
"""
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

import warnings; warnings.filterwarnings("ignore")
import json
import argparse
import numpy as np
import pandas as pd
import yfinance as yf

from scripts.improvement_harness import (
    fetch_df, build_triple_barrier_labels, walk_forward, compute_metrics,
    CURRENT_FEATURES, DEFAULT_STOCKS, EXTENDED_STOCKS,
    PASS_DIR_ACC, PASS_UP_PREC,
)
from models.predictor import _build_features

C24_FEATURES = [
    "djia_ret_1d",
    "djia_ret_5d",
    "djia_ma20_ratio",
    "vix_level",
    "vix_change_5d",
]

NEW_FEATURES = CURRENT_FEATURES + C24_FEATURES

_macro_cache: dict = {}


def _fetch_macro(start: str, end: str) -> tuple[pd.Series | None, pd.Series | None]:
    key = f"{start}:{end}"
    if key in _macro_cache:
        return _macro_cache[key]
    djia = vix = None
    try:
        raw = yf.download(["^DJI", "^VIX"], start=start, end=end,
                          auto_adjust=True, progress=False)
        if not raw.empty:
            close = raw["Close"] if "Close" in raw.columns else raw.get("close")
            if close is not None and not close.empty:
                close.index = pd.to_datetime(close.index).strftime("%Y-%m-%d")
                if "^DJI" in close.columns:
                    djia = close["^DJI"].dropna()
                if "^VIX" in close.columns:
                    vix = close["^VIX"].dropna()
    except Exception as exc:
        print(f"  [warn] macro fetch failed: {exc}")
    result = (djia, vix)
    _macro_cache[key] = result
    return result


def _add_macro_to_feat(feat: pd.DataFrame, df: pd.DataFrame,
                       djia: pd.Series | None, vix: pd.Series | None) -> pd.DataFrame:
    """Append C24 macro columns to feat (already built by _build_features)."""
    date_col = df["date"].astype(str) if "date" in df.columns else None
    nan_col  = pd.Series(np.nan, index=feat.index)

    if djia is not None and date_col is not None:
        d = djia.to_dict()
        aligned = date_col.map(d).astype(float)
        feat["djia_ret_1d"]     = aligned.pct_change(1).shift(1).values  # non-leaking
        feat["djia_ret_5d"]     = aligned.pct_change(5).values
        ma20                     = aligned.rolling(20, min_periods=20).mean()
        feat["djia_ma20_ratio"] = (aligned / ma20.replace(0, np.nan)).values
    else:
        feat["djia_ret_1d"] = feat["djia_ret_5d"] = feat["djia_ma20_ratio"] = 0.0

    if vix is not None and date_col is not None:
        v = vix.to_dict()
        aligned_v = date_col.map(v).astype(float)
        feat["vix_level"]     = aligned_v.values
        feat["vix_change_5d"] = aligned_v.diff(5).values
    else:
        feat["vix_level"] = feat["vix_change_5d"] = 0.0

    # forward-fill then neutral-fill
    for col in C24_FEATURES:
        feat[col] = pd.Series(feat[col], index=feat.index).ffill().bfill().fillna(0.0)

    return feat


def load_stock(stock_no: str):
    df = fetch_df(stock_no)
    if df is None or df.empty:
        return None, None, None

    start = str(df["date"].min()) if "date" in df.columns else "2020-01-01"
    end   = str(df["date"].max()) if "date" in df.columns else "2025-01-01"
    djia, vix = _fetch_macro(start, end)

    feat   = _build_features(df)
    feat   = _add_macro_to_feat(feat, df, djia, vix)
    close  = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
    labels = build_triple_barrier_labels(close)
    return feat, labels, df


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--extended", action="store_true")
    args = parser.parse_args()

    stocks = EXTENDED_STOCKS if args.extended else DEFAULT_STOCKS
    tag    = "12-stock" if args.extended else "5-stock"
    suffix = "_12stock" if args.extended else ""

    print(f"\n=== C24: DJIA + VIX macro features [{tag}] ===")
    print(f"  Baseline features: {len(CURRENT_FEATURES)}")
    print(f"  C24 new features:  {C24_FEATURES}\n")

    per_stock = {}
    base_results, c24_results = [], []

    for stock_no in stocks:
        print(f"  {stock_no}...", end=" ", flush=True)
        feat, labels, _ = load_stock(stock_no)
        if feat is None:
            print("no data")
            continue

        m_base = walk_forward(feat, labels, CURRENT_FEATURES)
        m_c24  = walk_forward(feat, labels, NEW_FEATURES)

        per_stock[stock_no] = {"baseline": m_base, "c24": m_c24}
        base_results.append(m_base)
        c24_results.append(m_c24)

        b_dir = m_base.get("dir_accuracy", float("nan"))
        b_up  = m_base.get("up_precision", float("nan"))
        c_dir = m_c24.get("dir_accuracy",  float("nan"))
        c_up  = m_c24.get("up_precision",  float("nan"))
        print(f"base dir={b_dir}% ↑prec={b_up}%  →  c24 dir={c_dir}% ↑prec={c_up}%")

    def _avg(results, key):
        vals = [m[key] for m in results if m and not np.isnan(m.get(key, float("nan")))]
        return round(float(np.mean(vals)), 1) if vals else float("nan")

    base_agg = {"dir_accuracy": _avg(base_results, "dir_accuracy"),
                "up_precision": _avg(base_results, "up_precision")}
    c24_agg  = {"dir_accuracy": _avg(c24_results,  "dir_accuracy"),
                "up_precision": _avg(c24_results,  "up_precision")}

    passed = c24_agg["dir_accuracy"] >= PASS_DIR_ACC and c24_agg["up_precision"] >= PASS_UP_PREC

    print(f"\n  Baseline avg: dir={base_agg['dir_accuracy']}%  ↑prec={base_agg['up_precision']}%")
    print(f"  C24 avg:      dir={c24_agg['dir_accuracy']}%   ↑prec={c24_agg['up_precision']}%")
    print(f"  Gate (dir≥{PASS_DIR_ACC}% AND ↑prec≥{PASS_UP_PREC}%): {'PASS ✓' if passed else 'FAIL ✗'}")

    result = {
        "experiment":   "C24",
        "description":  "DJIA + VIX macro features",
        "new_features": C24_FEATURES,
        "stocks":       stocks,
        "aggregate":    {"baseline": base_agg, "c24": c24_agg},
        "passed":       passed,
        "pass_gate":    {"dir_accuracy": PASS_DIR_ACC, "up_precision": PASS_UP_PREC},
        "per_stock":    per_stock,
    }

    out = ROOT / f"docs/c24_result{suffix}.json"
    out.parent.mkdir(exist_ok=True)
    out.write_text(json.dumps(result, indent=2))
    print(f"\n  Saved: {out}")
    return 0 if passed else 1


if __name__ == "__main__":
    sys.exit(main())