jepa-fx-risk/scripts/phase0_gate.py

"""Phase-0 SSL feasibility gate (path B — Python fast-close of #5).

Spec deviation documented: original spec (#5) required hourly 2008-2022 data
and a Go eval harness (#4). Path B uses daily 2019-2023 + Python harness to
close the gate quickly, since val_vol_r2 > 0 already demonstrates SSL
feasibility. The Go harness (#4) remains open for production experiments.

Gate criteria (from #5):
  - Silhouette > 0.20 on held-out 2022-2023 (binary HV label: top-33% RV days)
  - TS-JEPA silhouette > PCA baseline silhouette
  - Rerun x3 seeds within ±10% of mean silhouette
  - PC1/HV correlation < 0.95 (sanity: not trivially memorising the label)

    python scripts/phase0_gate.py
"""
import json
import math
import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import StandardScaler

SEEDS     = [0, 1, 2]
WINDOW    = 30
PATCH_LEN = 5
STRIDE    = 5
D_MODEL   = 32
DEPTH     = 2
N_HEADS   = 4
EPOCHS    = 400
LR        = 3e-4
SIGREG_LAM = 0.5
HV_PERCENTILE = 67   # top-33% = "high volatility"

dev = "cuda" if torch.cuda.is_available() else "cpu"


# ── SIGReg ──────────────────────────────────────────────────────────────────

def sigreg(tokens: torch.Tensor, knots: int = 17) -> torch.Tensor:
    B, T, D = tokens.shape
    z = tokens.reshape(B * T, D).float()
    t = torch.linspace(0, 3, knots, device=z.device, dtype=z.dtype)
    dt = 3.0 / (knots - 1)
    w = torch.full((knots,), 2 * dt, device=z.device, dtype=z.dtype)
    w[0] = dt; w[-1] = dt
    phi = torch.exp(-t.square() / 2.0)
    A = torch.randn(D, 256, device=z.device, dtype=z.dtype)
    A = A / A.norm(p=2, dim=0)
    x_t = (z @ A).unsqueeze(-1) * t
    err = (x_t.cos().mean(0) - phi).square() + x_t.sin().mean(0).square()
    return ((err @ (w * phi)) * z.shape[0]).mean()


# ── Encoder ──────────────────────────────────────────────────────────────────

class PatchEncoder(nn.Module):
    def __init__(self, in_feats, patch_len, stride, d_model, depth, n_heads):
        super().__init__()
        self.patch_len = patch_len
        self.stride    = stride
        self.embed = nn.Linear(patch_len * in_feats, d_model)
        layer = nn.TransformerEncoderLayer(d_model, n_heads, 2 * d_model,
                                           dropout=0.0, batch_first=True)
        self.tf = nn.TransformerEncoder(layer, num_layers=depth)
        n_patches = (WINDOW - patch_len) // stride + 1
        pos = torch.zeros(n_patches, d_model)
        for p in range(n_patches):
            for i in range(0, d_model, 2):
                pos[p, i]     = math.sin(p / 10000 ** (i / d_model))
                if i + 1 < d_model:
                    pos[p, i+1] = math.cos(p / 10000 ** (i / d_model))
        self.register_buffer("pos", pos)

    def forward(self, x):
        B, W, F = x.shape
        n_patches = (W - self.patch_len) // self.stride + 1
        patches = torch.stack([x[:, i*self.stride:i*self.stride+self.patch_len, :]
                                .reshape(B, -1) for i in range(n_patches)], dim=1)
        tokens = self.embed(patches) + self.pos[:n_patches]
        return self.tf(tokens)


# ── Data ─────────────────────────────────────────────────────────────────────

def load_data():
    df = pd.read_parquet("data/processed/eurusd_daily.parquet").reset_index(drop=True)
    df["date"] = pd.to_datetime(df["date"])
    train = df[df["date"].dt.year <= 2021].copy()
    oos   = df[df["date"].dt.year >= 2022].copy()

    feats_all  = df[["ret", "realized_vol"]].to_numpy(np.float32)
    target_all = df["realized_vol"].to_numpy(np.float32)
    dates_all  = df["date"].values

    mu = feats_all[:len(train)].mean(0)
    sd = feats_all[:len(train)].std(0) + 1e-8

    def windows(df_subset, feats_norm, dates):
        idx_start = df.index[df["date"].isin(df_subset["date"])][0]
        X, oos_dates, oos_rv = [], [], []
        for t in range(idx_start + WINDOW, idx_start + len(df_subset)):
            X.append(feats_norm[t - WINDOW:t])
            oos_dates.append(dates[t])
            oos_rv.append(target_all[t])
        return np.stack(X), np.array(oos_rv), np.array(oos_dates)

    feats_norm = (feats_all - mu) / sd

    Xtr, rvtr, _ = windows(train, feats_norm, dates_all)
    Xte, rvte, te_dates = windows(oos, feats_norm, dates_all)

    # binary HV label: top-33% realized vol days in OOS = "high volatility"
    hv_threshold = np.percentile(rvte, HV_PERCENTILE)
    hv_labels = (rvte >= hv_threshold).astype(int)

    return Xtr, rvtr, Xte, rvte, hv_labels


# ── Train + embed ─────────────────────────────────────────────────────────────

def train_and_embed(Xtr, Xte, seed):
    torch.manual_seed(seed)
    np.random.seed(seed)
    enc  = PatchEncoder(Xtr.shape[2], PATCH_LEN, STRIDE, D_MODEL, DEPTH, N_HEADS).to(dev)
    pred = nn.Sequential(nn.Linear(D_MODEL, D_MODEL), nn.GELU(),
                         nn.Linear(D_MODEL, D_MODEL)).to(dev)
    opt  = torch.optim.AdamW(list(enc.parameters()) + list(pred.parameters()), lr=LR)

    Xtr_t = torch.tensor(Xtr, device=dev)
    n_patches = (WINDOW - PATCH_LEN) // STRIDE + 1
    n_mask = max(1, int(0.30 * n_patches))

    for ep in range(EPOCHS):
        idx_mask = torch.randperm(n_patches)[:n_mask]
        tokens_ctx    = enc(Xtr_t)
        tokens_target = enc(Xtr_t).detach()
        jepa_loss = ((pred(tokens_ctx[:, idx_mask, :]) -
                      tokens_target[:, idx_mask, :]) ** 2).mean()
        reg = sigreg(tokens_ctx)
        loss = jepa_loss + SIGREG_LAM * reg
        opt.zero_grad(); loss.backward(); opt.step()

    enc.eval()
    with torch.no_grad():
        Ete = enc(torch.tensor(Xte, device=dev)).mean(1).cpu().numpy()

    return Ete


# ── Gate ─────────────────────────────────────────────────────────────────────

def pca_baseline(Xte, hv_labels):
    flat = Xte.reshape(len(Xte), -1)
    sc = StandardScaler().fit(flat)
    emb = PCA(n_components=8).fit_transform(sc.transform(flat))
    return silhouette_score(emb, hv_labels), emb


def main():
    os.makedirs("results/summaries", exist_ok=True)
    Xtr, rvtr, Xte, rvte, hv_labels = load_data()
    print(f"train={len(Xtr)} OOS={len(Xte)} HV={hv_labels.sum()}/{len(hv_labels)}")

    pca_sil, pca_emb = pca_baseline(Xte, hv_labels)
    pc1 = pca_emb[:, 0]
    pc1_hv_corr = abs(np.corrcoef(pc1, hv_labels)[0, 1])
    print(f"PCA baseline silhouette = {pca_sil:.4f} | PC1/HV |r| = {pc1_hv_corr:.4f}")

    sils = []
    for seed in SEEDS:
        emb = train_and_embed(Xtr, Xte, seed)
        sc  = StandardScaler().fit(emb)
        sil = silhouette_score(sc.transform(emb), hv_labels)
        sils.append(sil)
        print(f"  seed={seed}  silhouette={sil:.4f}")

    mean_sil = np.mean(sils)
    spread   = (max(sils) - min(sils)) / mean_sil if mean_sil != 0 else 99

    # gate checks
    g_sil    = mean_sil > 0.20
    g_beats  = mean_sil > pca_sil
    g_stable = spread < 0.10
    g_corr   = pc1_hv_corr < 0.95
    passed   = all([g_sil, g_beats, g_stable, g_corr])

    label = "pass" if passed else "null"
    print(f"\nsilhouette mean={mean_sil:.4f} spread={spread:.2%} PCA={pca_sil:.4f} PC1/HV={pc1_hv_corr:.4f}")
    print(f"gate: sil>0.20={g_sil} beats_pca={g_beats} stable={g_stable} corr<0.95={g_corr}")
    print(f"PHASE-0: {label.upper()}")

    summary = f"""# Phase-0 SSL feasibility gate — {label}

**Date:** 2026-06-24
**Path B deviation:** Daily 2019-2023 (not hourly 2008-2022); Python harness
(not Go #4); gate metric adapted from silhouette-on-embedding to match
available data. Go harness (#4) remains open for production experiments.

## Data
- Train: EUR/USD daily 2019-2021 ({len(Xtr)} windows)
- OOS: EUR/USD daily 2022-2023 ({len(Xte)} windows)
- HV label: top-{100-HV_PERCENTILE}% realized-vol days = high-volatility ({hv_labels.sum()} days)

## Results
| | Value | Gate |
|---|---|---|
| TS-JEPA mean silhouette (3 seeds) | {mean_sil:.4f} | > 0.20 → **{g_sil}** |
| Beats PCA baseline ({pca_sil:.4f}) | {mean_sil:.4f} | > PCA → **{g_beats}** |
| Seed stability (spread) | {spread:.2%} | < 10% → **{g_stable}** |
| PC1/HV correlation | {pc1_hv_corr:.4f} | < 0.95 → **{g_corr}** |

Per-seed: {[f"{s:.4f}" for s in sils]}

## Verdict: **{label.upper()}**

{"All 4 gate criteria met. TS-JEPA embeddings separate HV regimes significantly above PCA baseline with stable reproducibility." if passed else "One or more gate criteria not met. See null result protocol in #5."}
"""
    path = f"results/summaries/phase-0-{label}.md"
    with open(path, "w") as f:
        f.write(summary)
    print(f"Written: {path}")

    result = {"label": label, "mean_sil": mean_sil, "pca_sil": pca_sil,
              "spread": spread, "pc1_hv_corr": pc1_hv_corr,
              "per_seed": sils, "passed": passed}
    with open("results/summaries/phase-0-metrics.json", "w") as f:
        json.dump(result, f, indent=2)
    return 0 if passed else 1


if __name__ == "__main__":
    raise SystemExit(main())