jepa-fx-risk/tests/test_hepa.py

"""Failing tests for HEPA backbone + Phase-1 supervised head + HPO in train.py.

Run: cd ~/dev/AI/jepa-fx-risk && .venv/bin/python -m pytest tests/test_hepa.py -v
These tests define what the backbone and head must satisfy BEFORE implementation.
"""
import math
import os
import torch
import torch.nn as nn
import pytest

# ── Tests import the classes from train.py ────────────────────────────────────
# They will fail until train.py implements: CausalEncoder, HorizonPredictor, vicreg_loss


def _import(env_overrides=None):
    import importlib.util, sys
    saved = {}
    if env_overrides:
        for k, v in env_overrides.items():
            saved[k] = os.environ.get(k)
            os.environ[k] = str(v)
    # Force fresh module load (env vars must be read at import time)
    name = f"train_{id(env_overrides)}"
    spec = importlib.util.spec_from_file_location(name, "train.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    if env_overrides:
        for k, orig in saved.items():
            if orig is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = orig
    return mod


@pytest.fixture(scope="module")
def train_mod():
    return _import()


# 1. CausalEncoder exists and has correct output shape
def test_causal_encoder_shape(train_mod):
    enc = train_mod.CausalEncoder(n_channels=2, patch_len=10, d_model=32, n_heads=4, depth=1)
    x = torch.randn(4, 60, 2)
    tokens = enc(x)          # should return all tokens (B, N, D) for JEPA pretraining
    assert tokens.shape == (4, 6, 32), f"expected (4, 6, 32), got {tokens.shape}"


# 2. CausalEncoder is actually causal: earlier token outputs don't change when later inputs change
def test_causal_masking(train_mod):
    enc = train_mod.CausalEncoder(n_channels=2, patch_len=10, d_model=32, n_heads=4, depth=2)
    enc.eval()
    torch.manual_seed(0)
    x = torch.randn(1, 60, 2)
    x_perturbed = x.clone()
    # non-uniform noise (constant shift absorbed by per-patch LayerNorm; variance change is not)
    torch.manual_seed(99)
    x_perturbed[:, 30:, :] += torch.randn_like(x[:, 30:, :]) * 5.0

    with torch.no_grad():
        h1 = enc(x)
        h2 = enc(x_perturbed)

    # First 3 tokens must be identical (causal — don't see future patches)
    assert torch.allclose(h1[:, :3, :], h2[:, :3, :], atol=1e-5), \
        "causal masking broken: early tokens change when later input changes"
    # Last token should differ (it can see the perturbed patches)
    assert not torch.allclose(h1[:, -1, :], h2[:, -1, :], atol=1e-5), \
        "last token should differ when later input changes"


# 3. HorizonPredictor exists, takes (h, delta_t_float) → same shape as h
def test_horizon_predictor_shape(train_mod):
    pred = train_mod.HorizonPredictor(d_model=32)
    h = torch.randn(4, 32)
    dt = torch.tensor([1.0, 2.0, 3.0, 1.0])
    out = pred(h, dt)
    assert out.shape == (4, 32), f"expected (4, 32), got {out.shape}"


# 4. vicreg_loss is a scalar and backward doesn't error
def test_vicreg_loss_backward(train_mod):
    h_pred = torch.randn(8, 32, requires_grad=True)
    h_target = torch.randn(8, 32)
    loss = train_mod.vicreg_loss(h_pred, h_target, alpha=0.1)
    assert loss.shape == (), f"expected scalar, got {loss.shape}"
    loss.backward()
    assert h_pred.grad is not None


# 5. Full JEPA step: encode context, predict future, compute loss, backward
def test_jepa_step_end_to_end(train_mod):
    enc = train_mod.CausalEncoder(n_channels=2, patch_len=10, d_model=32, n_heads=4, depth=1)
    pred = train_mod.HorizonPredictor(d_model=32)
    opt = torch.optim.SGD(list(enc.parameters()) + list(pred.parameters()), lr=1e-3)

    x = torch.randn(4, 60, 2)
    tokens = enc(x)          # (4, 6, 32)
    c, dt = 2, 2             # context position 2, horizon 2
    h_ctx = tokens[:, c, :]
    h_tgt = tokens[:, c + dt, :].detach()
    h_hat = pred(h_ctx, torch.full((4,), float(dt)))
    loss = train_mod.vicreg_loss(h_hat, h_tgt, alpha=0.1)
    opt.zero_grad(); loss.backward(); opt.step()
    assert loss.item() < 100, "loss exploded"


# 6. build() returns year-based OOS split (2022-2023); hourly gives many more windows
def test_build_year_split(train_mod):
    (Xtr, ytr), (Xte, yte) = train_mod.build()
    assert Xtr.shape[1] == train_mod.WINDOW
    assert Xte.shape[1] == train_mod.WINDOW
    assert len(Xtr) > 0 and len(Xte) > 0
    # OOS: daily ≈ 600; hourly ≈ 17,000 (2 years × ~8,500 trading hours/year)
    assert len(Xte) > 400, f"OOS too small: {len(Xte)}"


# 7. hourly build gives > 10× more training windows than daily
def test_build_hourly_more_windows(train_mod):
    import os
    if not os.path.exists("data/processed/eurusd_hourly.parquet"):
        pytest.skip("eurusd_hourly.parquet not present — run data:prepare:hourly first")
    (Xtr, _), _ = train_mod.build()
    # Daily had ~877 train windows; hourly with 2008-2021 should have > 50,000
    assert len(Xtr) > 10_000, f"expected >10k hourly train windows, got {len(Xtr)}"


# ── Phase-1: supervised head ──────────────────────────────────────────────────

# 8. SupervisedHead exists and maps (B, D) → (B,)
def test_supervised_head_shape(train_mod):
    D = 128
    head = train_mod.SupervisedHead(D)
    x = torch.randn(16, D)
    out = head(x)
    assert out.shape == (16,), f"expected (16,), got {out.shape}"


# 9. SupervisedHead gradient flows (not frozen)
def test_supervised_head_backward(train_mod):
    head = train_mod.SupervisedHead(64)
    x = torch.randn(8, 64)
    loss = head(x).mean()
    loss.backward()
    for name, p in head.named_parameters():
        assert p.grad is not None, f"no grad on {name}"


# 10. Phase-1 beats linear on nonlinear synthetic signal
def test_phase1_beats_linear_on_nonlinear(train_mod):
    """MLP head should outperform ridge regression on data with nonlinear structure."""
    import numpy as np
    torch.manual_seed(0); np.random.seed(0)
    N, D = 1000, 32
    # target = |h|² (quadratic — linear can't fit well)
    Etr = np.random.randn(N, D).astype(np.float32)
    ytr = (Etr ** 2).sum(axis=1)
    Ete = np.random.randn(200, D).astype(np.float32)
    yte = (Ete ** 2).sum(axis=1)

    # Ridge baseline
    A = np.hstack([Etr, np.ones((N, 1))])
    w = np.linalg.solve(A.T @ A + 1e-3 * np.eye(A.shape[1]), A.T @ ytr)
    pred_lin = np.hstack([Ete, np.ones((200, 1))]) @ w
    r2_lin = float(1 - ((yte - pred_lin) ** 2).sum() / ((yte - yte.mean()) ** 2).sum())

    # MLP head
    head = train_mod.SupervisedHead(D)
    opt  = torch.optim.Adam(head.parameters(), lr=1e-2)
    Xtr_t = torch.tensor(Etr); ytr_t = torch.tensor(ytr)
    for _ in range(300):
        loss = nn.functional.mse_loss(head(Xtr_t), ytr_t)
        opt.zero_grad(); loss.backward(); opt.step()

    head.eval()
    with torch.no_grad():
        pred_mlp = head(torch.tensor(Ete)).numpy()
    r2_mlp = float(1 - ((yte - pred_mlp) ** 2).sum() / ((yte - yte.mean()) ** 2).sum())

    assert r2_mlp > r2_lin + 0.05, (
        f"MLP R²={r2_mlp:.3f} should beat ridge R²={r2_lin:.3f} by >0.05 on quadratic target"
    )


# 11. main() returns phase1_r2 in metrics.json (integration — needs real data)
def test_metrics_json_has_phase1_r2(train_mod):
    import json
    if not os.path.exists("metrics.json"):
        pytest.skip("metrics.json not present — run train.py first")
    with open("metrics.json") as f:
        m = json.load(f)
    assert "phase1_r2" in m, f"phase1_r2 missing from metrics.json: {list(m.keys())}"
    assert m["phase1_r2"] > m["val_vol_r2"], (
        f"MLP head phase1_r2={m['phase1_r2']:.4f} should beat linear probe "
        f"val_vol_r2={m['val_vol_r2']:.4f}"
    )


# ── HPO: env-var knob overrides ───────────────────────────────────────────────

# 12. JEPA_WINDOW env var overrides WINDOW at import time
def test_env_override_window():
    mod = _import({"JEPA_WINDOW": "48"})
    assert mod.WINDOW == 48, f"expected WINDOW=48, got {mod.WINDOW}"


# 13. JEPA_D_MODEL and JEPA_DEPTH env vars work
def test_env_override_d_model_depth():
    mod = _import({"JEPA_D_MODEL": "64", "JEPA_DEPTH": "4"})
    assert mod.D_MODEL == 64, f"expected D_MODEL=64, got {mod.D_MODEL}"
    assert mod.DEPTH == 4, f"expected DEPTH=4, got {mod.DEPTH}"


# 14. hpo_sweep.py exists and generates correct config list
def test_hpo_sweep_configs():
    import importlib.util
    sweep_path = "scripts/hpo_sweep.py"
    if not os.path.exists(sweep_path):
        pytest.fail(f"{sweep_path} not found — implement it")
    spec = importlib.util.spec_from_file_location("hpo_sweep", sweep_path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    cfgs = list(mod.configs())
    assert len(cfgs) > 0, "configs() returned empty list"
    # Every config must have at least D_MODEL, DEPTH, WINDOW keys
    required = {"JEPA_D_MODEL", "JEPA_DEPTH", "JEPA_WINDOW"}
    for cfg in cfgs:
        assert required.issubset(cfg.keys()), f"config missing required keys: {cfg}"


# ── Option B: joint encoder fine-tuning in phase-1 ───────────────────────────

# 15. PHASE1_JOINT and PHASE1_ENCODER_LR knobs exist at module level
def test_joint_phase1_knobs():
    mod = _import({"JEPA_PHASE1_JOINT": "1", "JEPA_PHASE1_ENCODER_LR": "1e-5"})
    assert hasattr(mod, "PHASE1_JOINT"), "PHASE1_JOINT knob missing from train.py"
    assert hasattr(mod, "PHASE1_ENCODER_LR"), "PHASE1_ENCODER_LR knob missing from train.py"
    assert mod.PHASE1_JOINT is True
    assert abs(mod.PHASE1_ENCODER_LR - 1e-5) < 1e-12


# 16. PHASE1_JOINT defaults to True (joint mode on by default)
def test_joint_phase1_default_on():
    mod = _import()
    assert hasattr(mod, "PHASE1_JOINT"), "PHASE1_JOINT knob missing"
    assert mod.PHASE1_JOINT is True, f"PHASE1_JOINT default should be True, got {mod.PHASE1_JOINT}"


# 17. JEPA_PHASE1_JOINT=0 disables joint (env override works)
def test_joint_phase1_can_disable():
    mod = _import({"JEPA_PHASE1_JOINT": "0"})
    assert mod.PHASE1_JOINT is False, f"expected False, got {mod.PHASE1_JOINT}"


# 18. Encoder receives non-zero gradients when joint-training with the head
def test_joint_encoder_grad_flows(train_mod):
    """Gradient must flow into encoder when using two-param-group joint optimizer."""
    import torch.nn.functional as F
    enc  = train_mod.CausalEncoder(n_channels=2, patch_len=8, d_model=16, n_heads=2, depth=1)
    head = train_mod.SupervisedHead(16)
    enc.train(); head.train()
    opt = torch.optim.Adam([
        {"params": head.parameters(), "lr": 1e-3},
        {"params": enc.parameters(), "lr": 1e-5},
    ], weight_decay=1e-4)
    # Tiny batch: 4 windows of length 16 (= 2 patches of patch_len=8)
    X = torch.randn(4, 16, 2)
    y = torch.randn(4)
    tokens = enc(X)           # (4, 2, 16)
    h      = tokens[:, -1, :] # (4, 16) — last token
    pred   = head(h)
    loss   = F.mse_loss(pred, y)
    loss.backward()
    enc_grads = [p.grad for p in enc.parameters() if p.grad is not None]
    assert len(enc_grads) > 0, "no encoder params received gradients"
    assert any(g.abs().max().item() > 0 for g in enc_grads), "all encoder grads are zero"