Warrior_EA/research/altdata/screen_wti.py

"""EIA x WTI screen: petroleum fundamentals + positioning vs forward range/direction.

Same harness as screen.py (as-of join on published, 3x3 MI, circular-shift null,
family-wise max bar, +/- controls, conditional-on-trailing-range column), on D1
bars resampled from the decoded XTIUSD M1 file. Run BEFORE any EIA feature is
wired into the EA - survivors only.

Usage:
    python -m altdata.screen_wti [--perms 499]
"""
import argparse

import numpy as np
import pandas as pd

from .common import DATA_ROOT
from .screen import _tercile, asof_join, cmi_3x3x3, mi_3x3, true_range

M1_FILE = DATA_ROOT.parent / "sqxbars" / "XTIUSD_the5ers_M1.npz"


def load_wti_d1() -> pd.DataFrame:
    z = np.load(M1_FILE)
    df = pd.DataFrame(z["bars"], columns=list(z["columns"]))
    df["time"] = pd.to_datetime(df["time"], unit="ms")
    df = df.set_index("time")
    d1 = pd.DataFrame({
        "open": df["open"].resample("1D").first(),
        "high": df["high"].resample("1D").max(),
        "low": df["low"].resample("1D").min(),
        "close": df["close"].resample("1D").last(),
    }).dropna()
    return d1.reset_index().rename(columns={"index": "time"})


def eia_feature(fname: str, kind: str) -> pd.DataFrame:
    f = pd.read_csv(DATA_ROOT / "eia" / fname, parse_dates=["observed", "published"])
    f = f.sort_values("observed").reset_index(drop=True)
    v = f["value"]
    if kind == "chg1":
        v = v.diff(1)
    elif kind == "chg4":
        v = v.diff(4)
    elif kind == "deseason":                      # level minus 52-week rolling mean
        v = v - v.rolling(52, min_periods=26).mean()
    name = fname.replace(".csv", "") + "_" + kind
    return pd.DataFrame({"published": f["published"], name: v})


def cot_features_wti() -> pd.DataFrame:
    c = pd.read_csv(DATA_ROOT / "cot" / "XTIUSD_cot.csv", parse_dates=["observed", "published"])
    oi = c["Open_Interest_All"].replace(0, np.nan)
    spec = c["M_Money_Positions_Long_All"] - c["M_Money_Positions_Short_All"]
    out = pd.DataFrame({"published": c["published"]})
    out["cot_spec_net"] = spec / oi
    out["cot_idx_1y"] = out["cot_spec_net"].rolling(52, min_periods=26).rank(pct=True)
    out["cot_chg_4w"] = out["cot_spec_net"].diff(4)
    return out


def fred_series(sid: str, name: str, diff: int) -> pd.DataFrame:
    f = pd.read_csv(DATA_ROOT / "fred" / f"{sid}.csv", parse_dates=["observed", "published"])
    v = f["value"].diff(diff) if diff else f["value"]
    return pd.DataFrame({"published": f["published"], name: v})


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--perms", type=int, default=499)
    ap.add_argument("--seed", type=int, default=7)
    args = ap.parse_args()
    rng = np.random.default_rng(args.seed)

    bars = load_wti_d1()
    tr = true_range(bars)
    atr = tr.ewm(alpha=1 / 14, min_periods=14).mean()
    trailing = tr.rolling(5).sum() / atr

    sources = [
        eia_feature("crude_stocks_ex_spr.csv", "chg1"),
        eia_feature("crude_stocks_ex_spr.csv", "deseason"),
        eia_feature("crude_field_production.csv", "chg4"),
        eia_feature("refinery_utilization_pct.csv", "level"),
        eia_feature("spr_stocks.csv", "chg4"),
        cot_features_wti(),
        fred_series("VIXCLS", "vix_chg5", 5),
        fred_series("DTWEXBGS", "usd_chg5", 5),
    ]
    feats = pd.DataFrame(index=bars.index)
    for src in sources:
        joined = asof_join(bars, src)
        for c in joined.columns:
            feats[c] = joined[c]
    feats["CTRL_trailing_range"] = trailing
    feats["CTRL_noise"] = rng.standard_normal(len(bars))

    for target in ("range", "dir"):
        if target == "range":
            y_raw = tr.shift(-1).rolling(5).sum().shift(-4) / atr
        else:
            y_raw = (bars["close"].shift(-5) - bars["close"]) / atr
        valid = y_raw.notna() & trailing.notna()
        y_bin = _tercile(y_raw[valid].to_numpy())
        z_bin = _tercile(trailing[valid].to_numpy())
        print(f"\n=== XTIUSD: {int(valid.sum())} D1 bars, target={target}, {args.perms} perms ===")
        results, null_max = [], np.zeros(args.perms)
        for name in feats.columns:
            x = feats.loc[valid, name].to_numpy(dtype=float)
            ok = ~np.isnan(x)
            if ok.sum() < 500:
                print(f"  {name:34s} skipped ({ok.sum()} rows)")
                continue
            xo, yo, zo = x[ok], y_bin[ok], z_bin[ok]
            n = len(xo)
            obs, obs_c = mi_3x3(xo, yo), cmi_3x3x3(xo, yo, z_bin[ok])
            null = np.empty(args.perms)
            null_c = np.empty(args.perms)
            for k in range(args.perms):
                xs = np.roll(xo, rng.integers(63, n - 63))
                null[k] = mi_3x3(xs, yo)
                null_c[k] = cmi_3x3x3(xs, yo, zo)
            p = (np.sum(null >= obs) + 1) / (args.perms + 1)
            p_c = (np.sum(null_c >= obs_c) + 1) / (args.perms + 1)
            results.append((name, obs, p, obs_c, p_c))
            null_max = np.maximum(null_max, null)
        fam = np.quantile(null_max, 0.95)
        print(f"  family-wise 5% bar over {len(results)} features: {fam:.5f}")
        for name, obs, p, obs_c, p_c in sorted(results, key=lambda r: -r[1]):
            mark = "FAMILY" if obs > fam else "      "
            print(f"  {name:34s} MI {obs:.5f} p={p:.4f} {mark} | MI|vol {obs_c:.5f} p={p_c:.4f}")


if __name__ == "__main__":
    main()