feat(hpo): env-var knob overrides + sweep script (18 configs)

- train.py knobs all readable from JEPA_* env vars (JEPA_WINDOW, JEPA_D_MODEL,
  JEPA_DEPTH, etc.) so hpo_sweep.py can override without touching source
- scripts/hpo_sweep.py: 3×2×3 grid over D_MODEL × DEPTH × WINDOW,
  logs to results/hpo/hpo_results.jsonl with leaderboard at end
- 3 new tests: env override correctness, configs() schema validation
- 19/19 tests pass

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Taskfile.yml

@@ -18,6 +18,26 @@ tasks:
     deps: [generate]
     cmds: [go test ./... -race]
 
+  data:fetch:
+    desc: "Download EUR/USD M1 from histdata (set YEARS env var)"
+    cmds: [.venv/bin/python scripts/fetch_data.py]
+  data:fetch:historical:
+    desc: "Download EUR/USD M1 2008-2018 from histdata"
+    cmds:
+      - YEARS=2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018 .venv/bin/python scripts/fetch_data.py
+  data:prepare:daily:
+    desc: "Rebuild eurusd_daily.parquet from all M1 zips"
+    cmds: [.venv/bin/python scripts/prepare_data.py]
+  data:prepare:hourly:
+    desc: "Build eurusd_hourly.parquet from all M1 zips"
+    cmds: [.venv/bin/python scripts/prepare_hourly.py]
+  data:prepare:all:
+    desc: "Build both daily and hourly parquets"
+    deps: [data:prepare:daily, data:prepare:hourly]
+  data:test:
+    desc: "Run Python data pipeline tests"
+    cmds: [.venv/bin/python -m pytest tests/test_prepare_hourly.py tests/test_hepa.py -v]
+
   eval:probe:
     desc: "Run linear-probe (val_vol_r2) on embeddings from metrics.json"
     cmds: [./bin/eval -metric probe]

cmd/eval/main.go

@@ -1,11 +1,19 @@
 // cmd/eval — CLI driver for the jepa-fx-risk evaluation harness.
-// Reads embeddings from a parquet/npy-style JSON export (embeddings.json)
-// and targets from eurusd_daily.parquet, then runs the requested metric.
 //
 //	./bin/eval -metric probe|silhouette|erank [-emb embeddings.json]
 //
-// embeddings.json format: {"embeddings": [[...], ...], "dates": ["2022-01-03", ...]}
+// embeddings.json format (from train.py EXPORT_EMBEDDINGS=1):
-// Generated by train.py when run with EXPORT_EMBEDDINGS=1.
+//
+//	{
+//	  "embeddings":         [[...], ...],  // OOS frozen embeddings
+//	  "realized_vol":       [...],         // OOS target (next-day RV)
+//	  "hv_label":           [...],         // binary HV label (top-33%)
+//	  "train_embeddings":   [[...], ...],  // train-set frozen embeddings
+//	  "train_realized_vol": [...]          // train-set RV targets
+//	}
+//
+// eval:probe standardises both sets using train statistics (no leakage).
+// Falls back to internal 70/30 split of OOS if train_embeddings absent.
 package main
 
 import (
@@ -26,34 +34,55 @@ func main() {
 
 	log := slog.New(slog.NewJSONHandler(os.Stdout, nil))
 
-	emb, labels, y, err := loadEmbeddings(*embFile)
+	d, err := readJSON(*embFile)
 	if err != nil {
 		log.Error("load embeddings", "err", err)
 		os.Exit(1)
 	}
-	log.Info("loaded", "n", len(emb), "dim", len(emb[0]), "metric", *metric)
+	log.Info("loaded", "oos", len(d.Embeddings), "dim", len(d.Embeddings[0]),
+		"train", len(d.TrainEmbeddings), "metric", *metric)
 
 	switch *metric {
 	case "probe":
-		r2 := eval.LinearProbe(emb, y, 1e-3)
+		var r2 float64
+		if len(d.TrainEmbeddings) > 0 {
+			// standardise both sets using train statistics to prevent leakage
+			trEmb, mu, sd := standardiseCompute(d.TrainEmbeddings)
+			oosEmb := applyStandardise(d.Embeddings, mu, sd)
+			r2 = eval.LinearProbeTrainTest(trEmb, d.TrainRealizedVol, oosEmb, d.RealizedVol, 1e-3)
+			log.Info("probe mode", "fit_on", "train_embeddings", "eval_on", "oos")
+		} else {
+			// fallback: internal 70/30 split of OOS embeddings
+			oosEmb, mu, sd := standardiseCompute(d.Embeddings)
+			n70 := int(float64(len(oosEmb)) * 0.7)
+			oos70 := applyStandardise(d.Embeddings[n70:], mu, sd)
+			r2 = eval.LinearProbeTrainTest(oosEmb[:n70], d.RealizedVol[:n70],
+				oos70, d.RealizedVol[n70:], 1e-3)
+			log.Info("probe mode", "fit_on", "oos[0:70%]", "eval_on", "oos[70%:]")
+		}
 		fmt.Printf(`{"metric":"val_vol_r2","value":%.6f}`+"\n", r2)
 		log.Info("linear probe", "val_vol_r2", fmt.Sprintf("%.4f", r2))
+
 	case "silhouette":
-		if labels == nil {
+		if len(d.HVLabel) == 0 {
-			log.Error("silhouette requires HV labels in embeddings.json")
+			log.Error("silhouette requires hv_label in embeddings.json")
 			os.Exit(1)
 		}
-		sil, err := eval.Silhouette(emb, labels)
+		oosEmb := standardise(d.Embeddings)
+		sil, err := eval.Silhouette(oosEmb, d.HVLabel)
 		if err != nil {
 			log.Error("silhouette", "err", err)
 			os.Exit(1)
 		}
 		fmt.Printf(`{"metric":"silhouette","value":%.6f}`+"\n", sil)
 		log.Info("silhouette", "score", fmt.Sprintf("%.4f", sil))
+
 	case "erank":
-		er := eval.EffectiveRank(emb)
+		oosEmb := standardise(d.Embeddings)
+		er := eval.EffectiveRank(oosEmb)
 		fmt.Printf(`{"metric":"effective_rank","value":%.6f}`+"\n", er)
 		log.Info("effective rank", "erank", fmt.Sprintf("%.2f", er))
+
 	default:
 		log.Error("unknown metric", "metric", *metric)
 		os.Exit(1)
@@ -61,32 +90,45 @@ func main() {
 }
 
 type embJSON struct {
-	Embeddings [][]float64 `json:"embeddings"`
+	Embeddings       [][]float64 `json:"embeddings"`
-	Dates      []string    `json:"dates"`
+	Dates            []string    `json:"dates"`
-	RealizedVol []float64  `json:"realized_vol"`
+	RealizedVol      []float64   `json:"realized_vol"`
-	HVLabel     []int      `json:"hv_label"`
+	HVLabel          []int       `json:"hv_label"`
+	TrainEmbeddings  [][]float64 `json:"train_embeddings"`
+	TrainRealizedVol []float64   `json:"train_realized_vol"`
 }
 
-func loadEmbeddings(path string) (emb [][]float64, labels []int, y []float64, err error) {
+func readJSON(path string) (*embJSON, error) {
 	f, err := os.Open(path)
 	if err != nil {
-		return nil, nil, nil, fmt.Errorf("open %s: %w", path, err)
+		return nil, fmt.Errorf("open %s: %w", path, err)
 	}
 	defer func() { _ = f.Close() }()
-
 	var d embJSON
 	if err := json.NewDecoder(f).Decode(&d); err != nil {
-		return nil, nil, nil, fmt.Errorf("decode: %w", err)
+		return nil, fmt.Errorf("decode: %w", err)
 	}
 	if len(d.Embeddings) == 0 {
-		return nil, nil, nil, fmt.Errorf("empty embeddings in %s", path)
+		return nil, fmt.Errorf("empty embeddings in %s", path)
 	}
+	return &d, nil
+}
 
-	// standardise embeddings (zero mean, unit std) per dimension
+// standardise centres + scales to zero mean / unit std; returns normalised rows.
-	n, dim := len(d.Embeddings), len(d.Embeddings[0])
+func standardise(rows [][]float64) [][]float64 {
+	out, _, _ := standardiseCompute(rows)
+	return out
+}
+
+// standardiseCompute centres + scales and returns (normalised, mu, sd) for reuse.
+func standardiseCompute(rows [][]float64) ([][]float64, []float64, []float64) {
+	if len(rows) == 0 {
+		return rows, nil, nil
+	}
+	n, dim := len(rows), len(rows[0])
 	mu := make([]float64, dim)
-	for _, row := range d.Embeddings {
+	for _, r := range rows {
-		for j, v := range row {
+		for j, v := range r {
 			mu[j] += v
 		}
 	}
@@ -94,8 +136,8 @@ func loadEmbeddings(path string) (emb [][]float64, labels []int, y []float64, er
 		mu[j] /= float64(n)
 	}
 	sd := make([]float64, dim)
-	for _, row := range d.Embeddings {
+	for _, r := range rows {
-		for j, v := range row {
+		for j, v := range r {
 			diff := v - mu[j]
 			sd[j] += diff * diff
 		}
@@ -103,16 +145,24 @@ func loadEmbeddings(path string) (emb [][]float64, labels []int, y []float64, er
 	for j := range sd {
 		sd[j] = math.Sqrt(sd[j]/float64(n)) + 1e-8
 	}
-	norm := make([][]float64, n)
+	out := make([][]float64, n)
-	for i, row := range d.Embeddings {
+	for i, r := range rows {
-		norm[i] = make([]float64, dim)
+		out[i] = make([]float64, dim)
-		for j, v := range row {
+		for j, v := range r {
-			norm[i][j] = (v - mu[j]) / sd[j]
+			out[i][j] = (v - mu[j]) / sd[j]
 		}
 	}
+	return out, mu, sd
+}
 
-	if len(d.HVLabel) > 0 {
+// applyStandardise normalises rows using pre-computed mu and sd.
-		labels = d.HVLabel
+func applyStandardise(rows [][]float64, mu, sd []float64) [][]float64 {
+	out := make([][]float64, len(rows))
+	for i, r := range rows {
+		out[i] = make([]float64, len(r))
+		for j, v := range r {
+			out[i][j] = (v - mu[j]) / sd[j]
+		}
 	}
-	return norm, labels, d.RealizedVol, nil
+	return out
 }

internal/eval/eval.go

@@ -10,6 +10,58 @@ import (
 	"math"
 )
 
+// LinearProbeTrainTest fits ridge regression on (trainEmb, trainY) and evaluates
+// on (testEmb, testY). Returns OOS R². Use this for proper held-out evaluation.
+func LinearProbeTrainTest(trainEmb [][]float64, trainY []float64,
+	testEmb [][]float64, testY []float64, lambda float64) float64 {
+	n := len(trainEmb)
+	if n == 0 || len(testEmb) == 0 {
+		return 0
+	}
+	d := len(trainEmb[0])
+	p := d + 1
+
+	A := make([][]float64, n)
+	for i, e := range trainEmb {
+		row := make([]float64, p)
+		copy(row, e)
+		row[d] = 1.0
+		A[i] = row
+	}
+	AtA := make([][]float64, p)
+	for i := range AtA {
+		AtA[i] = make([]float64, p)
+	}
+	Aty := make([]float64, p)
+	for i := 0; i < n; i++ {
+		for j := 0; j < p; j++ {
+			Aty[j] += A[i][j] * trainY[i]
+			for k := 0; k < p; k++ {
+				AtA[j][k] += A[i][j] * A[i][k]
+			}
+		}
+	}
+	for j := 0; j < p; j++ {
+		AtA[j][j] += lambda
+	}
+	w := solveCholesky(AtA, Aty)
+
+	yMean := mean(testY)
+	var ssRes, ssTot float64
+	for i, e := range testEmb {
+		row := make([]float64, p)
+		copy(row, e)
+		row[d] = 1.0
+		pred := dot(row, w)
+		ssRes += (testY[i] - pred) * (testY[i] - pred)
+		ssTot += (testY[i] - yMean) * (testY[i] - yMean)
+	}
+	if ssTot == 0 {
+		return 0
+	}
+	return 1 - ssRes/ssTot
+}
+
 // LinearProbe fits a ridge regression (closed-form) on (emb, y) with regularisation λ
 // and returns R² on the same data. Call with train embeddings; probe on held-out by
 // splitting before calling.

scripts/hpo_sweep.py

@@ -0,0 +1,97 @@
+"""HPO sweep for jepa-fx-risk HEPA backbone.
+
+Runs train.py with different JEPA_* env overrides, logs results to
+results/hpo/hpo_results.jsonl. Each config writes its metrics.json then
+the result is appended to the JSONL.
+
+Usage:
+    python scripts/hpo_sweep.py
+    python scripts/hpo_sweep.py --dry-run   # print configs, don't train
+"""
+import argparse
+import json
+import os
+import subprocess
+import sys
+from datetime import datetime
+from itertools import product
+from pathlib import Path
+
+# ── Search space ──────────────────────────────────────────────────────────────
+SEARCH_SPACE = {
+    "JEPA_D_MODEL": [64, 128, 256],
+    "JEPA_DEPTH":   [2, 4],
+    "JEPA_WINDOW":  [120, 240, 480],
+}
+# Fixed: PATCH_LEN=24 (1-day patches), N_HEADS=4, EPOCHS=300, PHASE1_EPOCHS=200
+
+PYTHON = str(Path(sys.executable))
+OUT_DIR = Path("results/hpo")
+
+
+def configs():
+    """Yield all configs as dicts of JEPA_* env overrides."""
+    keys = list(SEARCH_SPACE.keys())
+    for vals in product(*SEARCH_SPACE.values()):
+        yield dict(zip(keys, vals))
+
+
+def run_config(cfg: dict, metrics_path: str = "metrics.json") -> dict:
+    env = {**os.environ, **{k: str(v) for k, v in cfg.items()}}
+    result = subprocess.run(
+        [PYTHON, "train.py"],
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        return {"config": cfg, "error": result.stderr[-500:]}
+    stdout_last = result.stdout.strip().split("\n")[-1]
+    with open(metrics_path) as f:
+        m = json.load(f)
+    return {
+        "config": cfg,
+        "val_vol_r2":  m.get("val_vol_r2"),
+        "phase1_r2":   m.get("phase1_r2"),
+        "stdout_last": stdout_last,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    out_file = OUT_DIR / "hpo_results.jsonl"
+
+    all_cfgs = list(configs())
+    print(f"HPO sweep: {len(all_cfgs)} configs")
+    for i, cfg in enumerate(all_cfgs):
+        label = " ".join(f"{k.replace('JEPA_','')}={v}" for k, v in cfg.items())
+        print(f"\n[{i+1}/{len(all_cfgs)}] {label}")
+        if args.dry_run:
+            continue
+        ts = datetime.utcnow().isoformat()
+        row = run_config(cfg)
+        row["ts"] = ts
+        with open(out_file, "a") as f:
+            f.write(json.dumps(row) + "\n")
+        if "error" in row:
+            print(f"  ERROR: {row['error'][:200]}")
+        else:
+            print(f"  val_vol_r2={row['val_vol_r2']:.4f}  phase1_r2={row['phase1_r2']:.4f}")
+
+    if not args.dry_run:
+        # Print leaderboard
+        rows = [json.loads(l) for l in open(out_file) if l.strip()]
+        rows = [r for r in rows if "error" not in r]
+        rows.sort(key=lambda r: r.get("phase1_r2", -999), reverse=True)
+        print("\n── Leaderboard (by phase1_r2) ─────────────────────────")
+        for r in rows[:5]:
+            cfg_str = " ".join(f"{k.replace('JEPA_','')}={v}" for k,v in r["config"].items())
+            print(f"  {r['phase1_r2']:.4f}  {cfg_str}")
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_hourly.py

@@ -0,0 +1,108 @@
+"""Prepare EUR/USD hourly OHLCV + realized vol from histdata M1 zips.
+
+Aggregates all M1 bars in data/raw/DAT_ASCII_EURUSD_M1_*.zip to hourly.
+Realized vol per hour = sqrt(sum(log-return²)) over the constituent M1 bars.
+Weekend hours are naturally absent (FX market closed Sat/Sun); NO interpolation.
+Hours with fewer than MIN_BARS M1 bars are dropped (holidays, thin sessions).
+
+Output: data/processed/eurusd_hourly.parquet
+Columns: datetime (UTC, tz-naive), close, ret (log), realized_vol
+
+    python scripts/prepare_hourly.py
+    RAW=data/raw OUT=data/processed/eurusd_hourly.parquet python scripts/prepare_hourly.py
+"""
+import glob
+import os
+import zipfile
+
+import numpy as np
+import pandas as pd
+
+RAW_DEFAULT = "data/raw"
+OUT_DEFAULT = "data/processed/eurusd_hourly.parquet"
+MIN_BARS    = 30   # drop hours thinner than this (holidays, DST boundary artefacts)
+
+
+# ── Core transformation ──────────────────────────────────────────────────────
+
+def resample_to_hourly(m1: pd.DataFrame) -> pd.DataFrame:
+    """Aggregate M1 DataFrame to hourly bars.
+
+    Args:
+        m1: DataFrame with columns ['ts' (datetime), 'close' (float)]
+
+    Returns:
+        DataFrame with columns ['datetime', 'close', 'ret', 'realized_vol']
+        sorted by datetime; hours with fewer than MIN_BARS M1 ticks dropped.
+    """
+    m1 = m1.sort_values("ts").copy()
+    m1["log_r"] = np.log(m1["close"]).diff()
+    m1["hour"]  = m1["ts"].dt.floor("h")
+
+    agg = m1.groupby("hour").agg(
+        close       = ("close",  "last"),
+        realized_vol= ("log_r",  lambda x: np.sqrt(np.nansum(x.values ** 2))),
+        n_bars      = ("log_r",  "count"),
+    ).reset_index()
+
+    agg = agg[agg["n_bars"] >= MIN_BARS].copy()
+    agg["ret"] = np.log(agg["close"]).diff()
+    agg = agg.dropna(subset=["ret"]).reset_index(drop=True)
+    agg = agg.rename(columns={"hour": "datetime"})
+    return agg[["datetime", "close", "ret", "realized_vol"]]
+
+
+def load_m1_from_zips(raw_dir: str) -> pd.DataFrame:
+    """Load and concatenate all M1 zips from raw_dir (histdata format)."""
+    pattern = os.path.join(raw_dir, "DAT_ASCII_EURUSD_M1_*.zip")
+    zips = sorted(glob.glob(pattern))
+    if not zips:
+        raise FileNotFoundError(f"No M1 zips found at {pattern}")
+    frames = []
+    for zp in zips:
+        with zipfile.ZipFile(zp) as z:
+            csv = [n for n in z.namelist() if n.endswith(".csv")][0]
+            with z.open(csv) as f:
+                df = pd.read_csv(
+                    f, sep=";", header=None,
+                    names=["dt", "open", "high", "low", "close", "vol"],
+                )
+        df["ts"] = pd.to_datetime(df["dt"], format="%Y%m%d %H%M%S")
+        frames.append(df[["ts", "close"]])
+        print(f"  loaded {os.path.basename(zp)}: {len(df):,} rows")
+    return pd.concat(frames).sort_values("ts").reset_index(drop=True)
+
+
+def build_hourly_parquet(
+    raw_dir: str = RAW_DEFAULT,
+    out_path: str = OUT_DEFAULT,
+) -> pd.DataFrame:
+    """Full pipeline: load all M1 zips → hourly parquet. Returns the DataFrame."""
+    print(f"Loading M1 zips from {raw_dir}...")
+    m1 = load_m1_from_zips(raw_dir)
+    print(f"Total M1 bars: {len(m1):,}  ({m1['ts'].min().date()} → {m1['ts'].max().date()})")
+
+    print("Resampling to hourly...")
+    hourly = resample_to_hourly(m1)
+    print(f"Hourly rows: {len(hourly):,}  ({hourly['datetime'].min()} → {hourly['datetime'].max()})")
+
+    # Sanity: COVID crash (Mar 2020) should show realized vol spike if data covers it
+    if hourly["datetime"].dt.year.isin([2020]).any():
+        rv = hourly.set_index("datetime")["realized_vol"]
+        try:
+            mar20 = rv["2020-03-01":"2020-03-31"].max()
+            typ   = rv["2019-01-01":"2019-12-31"].median()
+            print(f"Sanity — median 2019 RV: {typ:.6f} | max Mar-2020 RV: {mar20:.6f} | spike ×{mar20/typ:.1f}")
+        except Exception:
+            pass
+
+    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
+    hourly.to_parquet(out_path, index=False)
+    print(f"Written: {out_path}")
+    return hourly
+
+
+if __name__ == "__main__":
+    raw_dir  = os.environ.get("RAW",  RAW_DEFAULT)
+    out_path = os.environ.get("OUT",  OUT_DEFAULT)
+    build_hourly_parquet(raw_dir=raw_dir, out_path=out_path)

tests/test_hepa.py

@@ -0,0 +1,229 @@
+"""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}"

tests/test_prepare_hourly.py

@@ -0,0 +1,126 @@
+"""Failing tests for scripts/prepare_hourly.py.
+
+Tests the M1 → hourly aggregation logic using synthetic data before touching
+real downloads.
+
+Run: cd ~/dev/AI/jepa-fx-risk && .venv/bin/python -m pytest tests/test_prepare_hourly.py -v
+"""
+import numpy as np
+import pandas as pd
+import pytest
+import importlib.util, sys, os
+
+
+def _import():
+    spec = importlib.util.spec_from_file_location(
+        "prepare_hourly", "scripts/prepare_hourly.py"
+    )
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+@pytest.fixture(scope="module")
+def ph():
+    return _import()
+
+
+def _make_m1(n_days: int = 3, price: float = 1.1000, noise: float = 0.0005) -> pd.DataFrame:
+    """Synthetic M1 DataFrame starting 2020-01-06 (Monday), 390 ticks/day."""
+    rng = np.random.default_rng(42)
+    # generate full trading hours: Mon-Fri 00:00-23:59 (FX is 24h weekday)
+    start = pd.Timestamp("2020-01-06 00:00:00")  # Monday
+    periods = n_days * 24 * 60
+    ts = pd.date_range(start, periods=periods, freq="min")
+    # remove weekends
+    ts = ts[ts.day_of_week < 5]
+    prices = price + np.cumsum(rng.normal(0, noise, len(ts)))
+    return pd.DataFrame({"ts": ts, "close": prices})
+
+
+# 1. resample_to_hourly: DataFrame has correct columns
+def test_columns(ph):
+    m1 = _make_m1()
+    hourly = ph.resample_to_hourly(m1)
+    assert set(["datetime", "close", "ret", "realized_vol"]).issubset(hourly.columns), \
+        f"missing columns: {hourly.columns.tolist()}"
+
+
+# 2. No cross-weekend interpolation: gap between Friday 23:xx and Sunday/Monday must remain
+def test_no_weekend_interpolation(ph):
+    # Make 2 days: Friday + Monday (skip Saturday/Sunday)
+    fri = pd.date_range("2020-01-10 00:00", "2020-01-10 23:59", freq="min")  # Friday
+    mon = pd.date_range("2020-01-13 00:00", "2020-01-13 23:59", freq="min")  # Monday
+    ts = fri.append(mon)
+    prices = 1.1 + np.cumsum(np.random.default_rng(0).normal(0, 0.0001, len(ts)))
+    m1 = pd.DataFrame({"ts": ts, "close": prices})
+    hourly = ph.resample_to_hourly(m1)
+    dates = pd.DatetimeIndex(hourly["datetime"]).date
+    import datetime
+    sat = datetime.date(2020, 1, 11)
+    sun = datetime.date(2020, 1, 12)
+    assert sat not in dates and sun not in dates, "weekend rows found in hourly output"
+
+
+# 3. Realized vol = sqrt(sum(r²)) over minute returns in each hour
+def test_realized_vol_formula(ph):
+    # Two hours: anchor gives 10:00 a valid ret; measurement hour has one known log-return.
+    ts0 = pd.date_range("2020-01-06 09:00", periods=60, freq="min")
+    ts1 = pd.date_range("2020-01-06 10:00", periods=60, freq="min")
+    prices0 = np.ones(60) * 1.0
+    # price jumps at minute 1 and STAYS (no reversion) → one non-zero log-return
+    prices1 = np.full(60, np.exp(0.01))
+    prices1[0] = 1.0  # only first tick is at 1.0; jump happens at tick 1
+    m1 = pd.DataFrame({
+        "ts": np.concatenate([ts0, ts1]),
+        "close": np.concatenate([prices0, prices1]),
+    })
+    hourly = ph.resample_to_hourly(m1)
+    assert len(hourly) >= 1, "no rows after resample"
+    rv = hourly.iloc[-1]["realized_vol"]
+    expected = np.sqrt(0.01 ** 2)
+    assert abs(rv - expected) < 1e-6, f"realized_vol={rv:.8f}, expected≈{expected:.8f}"
+
+
+# 4. Only hours with ≥ 30 M1 bars are kept (thin hours dropped)
+def test_thin_hours_dropped(ph):
+    # 4 hours: pre-anchor gives 09:00 a valid ret; full survives; thin (11:00) is dropped.
+    # pre-anchor (08:00): gives 09:00 a valid ret
+    # anchor (09:00): 60 bars, valid ret  → kept
+    # full  (10:00): 60 bars, valid ret  → kept
+    # thin  (11:00): 10 bars             → dropped
+    # Result: 3 hourly candidates, first (pre-anchor) gets NaN ret → dropped → 2 rows
+    pre   = pd.date_range("2020-01-06 08:00", periods=60, freq="min")
+    anchor= pd.date_range("2020-01-06 09:00", periods=60, freq="min")
+    full  = pd.date_range("2020-01-06 10:00", periods=60, freq="min")
+    thin  = pd.date_range("2020-01-06 11:00", periods=10, freq="min")
+    ts = pre.append(anchor).append(full).append(thin)
+    m1 = pd.DataFrame({"ts": ts, "close": np.ones(len(ts)) * 1.1})
+    hourly = ph.resample_to_hourly(m1)
+    assert len(hourly) == 2, f"expected 2 rows (pre-anchor NaN ret dropped + thin dropped), got {len(hourly)}"
+
+
+# 5. Output parquet path and schema (integration — reads actual M1 zips if present)
+def test_output_schema_from_zips(ph, tmp_path):
+    # Build a minimal fake zip structure
+    import zipfile, io
+    # synthetic M1 CSV (histdata format: YYYYMMDD HHMMSS;O;H;L;C;V)
+    rows = []
+    for h in range(24):
+        for m in range(60):
+            rows.append(f"20200106 {h:02d}{m:02d}00;1.10000;1.10100;1.09900;1.10000;100")
+    csv_content = "\n".join(rows).encode()
+    zip_buf = io.BytesIO()
+    with zipfile.ZipFile(zip_buf, "w") as zf:
+        zf.writestr("DAT_ASCII_EURUSD_M1_2020.csv", csv_content)
+    zip_buf.seek(0)
+    raw_dir = tmp_path / "raw"
+    raw_dir.mkdir()
+    (raw_dir / "DAT_ASCII_EURUSD_M1_2020.zip").write_bytes(zip_buf.read())
+
+    out_path = str(tmp_path / "eurusd_hourly.parquet")
+    ph.build_hourly_parquet(raw_dir=str(raw_dir), out_path=out_path)
+    assert os.path.exists(out_path), "output parquet not created"
+    df = pd.read_parquet(out_path)
+    assert set(["datetime", "close", "ret", "realized_vol"]).issubset(df.columns)
+    assert len(df) > 0

train.py

@@ -1,14 +1,13 @@
 """train.py — autoresearch agent file (only this may be edited).
 
-TS-JEPA backbone with SIGReg regularization (Balestriero & LeCun, LeJEPA
+HEPA backbone (Petersen et al., arXiv:2605.11130, ICML 2026 Spotlight):
-arXiv:2511.08544; time-series placement from ChronoJEPA arXiv: 2505.XXXXX).
+Causal Transformer pre-trained via horizon-conditioned JEPA. Predictor
+maps (h_t, Δt) → predicted future embedding; loss = VICReg (L1 alignment
+on L2-normalised reps + variance-covariance regulariser, no stop-gradient).
+Probe: ridge regression on the last-token embedding (true OOS split).
 
-PatchTST-style encoder over windowed daily [return, realized_vol] → FREEZE →
+Agent may tune: encoder depth/width, patch geometry, ALPHA, DELTA_T_MAX,
-linear probe predicts NEXT-day realized vol → val_vol_r2 (OOS R²).
+optimizer, LR. Do NOT touch prepare_data.py, loop.py, or the data pipeline.
-Writes metrics.json — the single scalar the loop reads.
-
-Agent may tune: encoder depth/width, patch geometry, mask strategy, SIGReg
-lambda, optimizer. Do NOT touch prepare_data.py, loop.py, or the data pipeline.
 """
 import json
 import math
@@ -16,19 +15,24 @@ import numpy as np
 import pandas as pd
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
-# --- agent-tunable knobs ---
+# --- agent-tunable knobs (all overridable via JEPA_* env vars for HPO) ---
-WINDOW     = 60     # INCREASED lookback for better volatility persistence capture
+import os as _os
-PATCH_LEN  = 5      # time-patch size (must divide WINDOW)
+USE_HOURLY    = True
-STRIDE     = 5
+WINDOW        = int(_os.environ.get("JEPA_WINDOW",        240))
-D_MODEL    = 64     # transformer hidden dim - INCREASED for capacity
+PATCH_LEN     = int(_os.environ.get("JEPA_PATCH_LEN",     24))
-DEPTH      = 2      # transformer layers
+D_MODEL       = int(_os.environ.get("JEPA_D_MODEL",       128))
-N_HEADS    = 4
+DEPTH         = int(_os.environ.get("JEPA_DEPTH",         2))
-MASK_FRAC  = 0.50   # INCREASED mask fraction to force the encoder to learn better global representations
+N_HEADS       = int(_os.environ.get("JEPA_N_HEADS",       4))
-SIGREG_LAM = 0.01   # SIGReg weight (λ) - REDUCED to allow more representation capacity
+ALPHA         = float(_os.environ.get("JEPA_ALPHA",       0.1))
-EPOCHS     = 300
+DELTA_T_MAX   = int(_os.environ.get("JEPA_DELTA_T_MAX",   3))
-LR         = 3e-4
+BATCH_SIZE    = int(_os.environ.get("JEPA_BATCH_SIZE",    512))
-SEED       = 0
+EPOCHS        = int(_os.environ.get("JEPA_EPOCHS",        300))
+LR            = float(_os.environ.get("JEPA_LR",          3e-4))
+PHASE1_EPOCHS = int(_os.environ.get("JEPA_PHASE1_EPOCHS", 200))
+PHASE1_LR     = float(_os.environ.get("JEPA_PHASE1_LR",   1e-3))
+SEED          = int(_os.environ.get("JEPA_SEED",          0))
 # ---------------------------
 
 torch.manual_seed(SEED)
@@ -36,139 +40,266 @@ np.random.seed(SEED)
 dev = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-# ── SIGReg (from LeJEPA/ChronoJEPA, token-level placement) ─────────────────
+# ── VICReg pretraining loss ──────────────────────────────────────────────────
 
-def sigreg(tokens: torch.Tensor, knots: int = 17) -> torch.Tensor:
+def vicreg_loss(h_pred: torch.Tensor, h_target: torch.Tensor, alpha: float = 0.1) -> torch.Tensor:
-    """Epps-Pulley test statistic pushes token embeddings toward isotropic Gaussian.
+    """L = (1-α)·L1(normalize(ĥ), normalize(h*)) + α·(L_var + L_cov).
 
-    tokens: (B, T, D) — applied per-token, averaged across B and T.
+    Both encoders receive gradients (joint training — no stop-grad on h_target).
+    Variance-covariance terms prevent embedding collapse.
     """
-    B, T, D = tokens.shape
+    pred_n = F.normalize(h_pred, dim=-1)
-    z = tokens.reshape(B * T, D)           # (N, D)
+    targ_n = F.normalize(h_target, dim=-1)
-    t = torch.linspace(0, 3, knots, device=z.device, dtype=z.float().dtype)
+    l1 = F.l1_loss(pred_n, targ_n)
-    dt = 3.0 / (knots - 1)
+    # variance hinge: push each feature std toward ≥ 1
-    w = torch.full((knots,), 2 * dt, device=z.device, dtype=z.float().dtype)
+    std = h_pred.std(dim=0) + 1e-4
-    w[0] = dt; w[-1] = dt
+    l_var = F.relu(1.0 - std).mean()
-    phi = torch.exp(-t.square() / 2.0)
+    # covariance penalty: decorrelate features
-
+    B, D = h_pred.shape
-    A = torch.randn(D, 256, device=z.device, dtype=z.float().dtype)
+    h_c = h_pred - h_pred.mean(dim=0, keepdim=True)
-    A = A / A.norm(p=2, dim=0)
+    cov = (h_c.t() @ h_c) / max(B - 1, 1)
-    x_t = (z.float() @ A).unsqueeze(-1) * t              # (N, 256, knots)
+    off = cov - torch.diag(torch.diag(cov))
-    err = (x_t.cos().mean(0) - phi).square() + x_t.sin().mean(0).square()
+    l_cov = (off ** 2).sum() / D
-    return ((err @ (w * phi)) * z.shape[0]).mean()
+    return (1 - alpha) * l1 + alpha * (l_var + l_cov)
 
 
-# ── Encoder + Predictor ─────────────────────────────────────────────────────
+# ── CausalEncoder ─────────────────────────────────────────────────────────────
 
-class PatchEncoder(nn.Module):
+class CausalEncoder(nn.Module):
-    """PatchTST-style encoder for univariate windows."""
+    """Non-overlapping patches → per-patch LayerNorm → causal Transformer → all tokens (B, N, D).
-    def __init__(self, in_feats, patch_len, stride, d_model, depth, n_heads):
+
+    Per-patch LayerNorm instead of full-window RevIN: each patch is normalised
+    using only its own timesteps, so no future statistics leak into past tokens.
+    Use [:, -1, :] for probing (last token sees full context).
+    Use [:, c, :] for JEPA pretraining (context-at-c).
+    """
+    def __init__(self, n_channels: int, patch_len: int, d_model: int,
+                 n_heads: int, depth: int):
         super().__init__()
         self.patch_len = patch_len
-        self.stride    = stride
         self.d_model   = d_model
-        self.embed = nn.Linear(patch_len * in_feats, d_model)
+        patch_dim = patch_len * n_channels
+        self.patch_norm = nn.LayerNorm(patch_dim)   # applied per-patch, no future leakage
+        self.embed     = nn.Linear(patch_dim, 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)
+        self.tf   = nn.TransformerEncoder(layer, num_layers=depth)
-        n_patches = (WINDOW - patch_len) // stride + 1
+        self.norm = nn.LayerNorm(d_model)
-        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: torch.Tensor) -> torch.Tensor:
-        # x: (B, W, F)  → patches → (B, T, D)
         B, W, F = x.shape
-        n_patches = (W - self.patch_len) // self.stride + 1
+        P = self.patch_len
-        patches = torch.stack([x[:, i*self.stride:i*self.stride+self.patch_len, :]
+        N = W // P
-                                .reshape(B, -1) for i in range(n_patches)], dim=1)
+        tokens = x[:, :N * P, :].reshape(B, N, P * F)
-        tokens = self.embed(patches) + self.pos[:n_patches]
+        tokens = self.embed(self.patch_norm(tokens))
-        return self.tf(tokens)  # (B, T, D)
+        # sinusoidal PE
+        pos = torch.arange(N, device=x.device).float()
+        div = torch.exp(torch.arange(0, self.d_model, 2, device=x.device).float()
+                        * -(math.log(10000.0) / self.d_model))
+        pe = torch.zeros(N, self.d_model, device=x.device)
+        pe[:, 0::2] = torch.sin(pos.unsqueeze(1) * div)
+        pe[:, 1::2] = torch.cos(pos.unsqueeze(1) * div)
+        tokens = tokens + pe
+        # causal mask
+        mask = nn.Transformer.generate_square_subsequent_mask(N, device=x.device)
+        return self.norm(self.tf(tokens, mask=mask, is_causal=True))
 
 
-class Predictor(nn.Module):
+# ── HorizonPredictor ─────────────────────────────────────────────────────────
-    def __init__(self, d_model):
+
+class HorizonPredictor(nn.Module):
+    """MLP(cat(h_t, Δt)) → predicted future embedding."""
+    def __init__(self, d_model: int):
         super().__init__()
-        self.net = nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(),
+        self.net = nn.Sequential(
-                                 nn.Linear(d_model, d_model))
+            nn.Linear(d_model + 1, d_model), nn.GELU(),
-    def forward(self, x):
+            nn.Linear(d_model, d_model),     nn.GELU(),
-        return self.net(x)
+            nn.Linear(d_model, d_model),
+        )
+
+    def forward(self, h: torch.Tensor, delta_t: torch.Tensor) -> torch.Tensor:
+        dt = delta_t.float().unsqueeze(-1)
+        return self.net(torch.cat([h, dt], dim=-1))
 
 
-# ── Data ────────────────────────────────────────────────────────────────────
+# ── Phase-1 supervised head ──────────────────────────────────────────────────
+
+class SupervisedHead(nn.Module):
+    """Small MLP trained on frozen HEPA embeddings to predict next-period realized vol."""
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, d_model // 2), nn.GELU(),
+            nn.Linear(d_model // 2, 1),
+        )
+
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        return self.net(h).squeeze(-1)
+
+
+# ── Data ─────────────────────────────────────────────────────────────────────
 
 def build():
-    df = pd.read_parquet("data/processed/eurusd_daily.parquet").reset_index(drop=True)
+    """Year-based split: encoder trains on ≤2021; probe evaluates on ≥2022 OOS.
+
+    Uses eurusd_hourly.parquet when USE_HOURLY=True and the file exists;
+    falls back to eurusd_daily.parquet otherwise.
+    """
+    import os
+    hourly_path = "data/processed/eurusd_hourly.parquet"
+    daily_path  = "data/processed/eurusd_daily.parquet"
+    if USE_HOURLY and os.path.exists(hourly_path):
+        df = pd.read_parquet(hourly_path).reset_index(drop=True)
+        df["date"] = pd.to_datetime(df["datetime"])
+    else:
+        df = pd.read_parquet(daily_path).reset_index(drop=True)
+        df["date"] = pd.to_datetime(df["date"])
     feats  = df[["ret", "realized_vol"]].to_numpy(np.float32)
     target = df["realized_vol"].to_numpy(np.float32)
-    X, y = [], []
+    tr_idx = df.index[df["date"].dt.year <= 2021].tolist()
-    for t in range(WINDOW, len(df) - 1):
+    te_idx = df.index[df["date"].dt.year >= 2022].tolist()
-        X.append(feats[t - WINDOW:t])
+    mu = feats[:tr_idx[-1]+1].mean(0)
-        y.append(target[t + 1])
+    sd = feats[:tr_idx[-1]+1].std(0) + 1e-8
-    X = np.stack(X); y = np.array(y, np.float32)
+    fn = (feats - mu) / sd
-    n_tr = int(0.7 * len(X))
+    def windows(idx):
-    mu  = X[:n_tr].mean((0, 1))
+        X, y = [], []
-    sd  = X[:n_tr].std((0, 1)) + 1e-8
+        for t in idx:
-    X   = (X - mu) / sd
+            if t - WINDOW >= 0 and t + 1 < len(df):
-    return (X[:n_tr], y[:n_tr]), (X[n_tr:], y[n_tr:])
+                X.append(fn[t - WINDOW:t]); y.append(target[t + 1])
+        return np.stack(X).astype(np.float32), np.array(y, np.float32)
+    return windows(tr_idx), windows(te_idx)
 
 
-# ── Training ─────────────────────────────────────────────────────────────────
+# ── Training ──────────────────────────────────────────────────────────────────
 
 def main():
     (Xtr, ytr), (Xte, yte) = build()
-    n_feats  = Xtr.shape[2]
+    n_feats   = Xtr.shape[2]
-    Xtr_t    = torch.tensor(Xtr, device=dev)
+    n_patches = WINDOW // PATCH_LEN
-    enc      = PatchEncoder(n_feats, PATCH_LEN, STRIDE, D_MODEL, DEPTH, N_HEADS).to(dev)
+    N_tr      = len(Xtr)
-    pred     = Predictor(D_MODEL).to(dev)
+    bs        = min(BATCH_SIZE, N_tr)
-    opt      = torch.optim.AdamW(list(enc.parameters()) + list(pred.parameters()), lr=LR)
 
-    n_patches = (WINDOW - PATCH_LEN) // STRIDE + 1
+    enc  = CausalEncoder(n_feats, PATCH_LEN, D_MODEL, N_HEADS, DEPTH).to(dev)
-    n_mask    = max(1, int(MASK_FRAC * n_patches))
+    pred = HorizonPredictor(D_MODEL).to(dev)
+    opt  = torch.optim.AdamW(list(enc.parameters()) + list(pred.parameters()), lr=LR)
 
     for ep in range(EPOCHS):
-        # JEPA: predict masked-out patch tokens from visible tokens
+        # Random mini-batch (avoids OOM on large hourly dataset)
-        idx_mask = torch.randperm(n_patches)[:n_mask]
+        idx_b = torch.randperm(N_tr)[:bs]
-        ctx_mask = torch.ones(n_patches, dtype=torch.bool, device=dev)
+        Xb    = torch.tensor(Xtr[idx_b.numpy()], device=dev)
-        ctx_mask[idx_mask] = False
 
-        tokens_ctx    = enc(Xtr_t)                 # encode all (B, T, D)
+        # Sample random context position and horizon
-        tokens_target = enc(Xtr_t).detach()        # target (frozen): same input, no grad
+        c  = torch.randint(0, n_patches - 1, ()).item()
-        pred_out = pred(tokens_ctx[:, idx_mask, :])
+        dt = torch.randint(1, max(2, min(DELTA_T_MAX, n_patches - 1 - c) + 1), ()).item()
-        jepa_loss = ((pred_out - tokens_target[:, idx_mask, :]) ** 2).mean()
+
-        reg_loss  = sigreg(tokens_ctx)
+        tokens = enc(Xb)                        # (bs, N, D)
-        loss      = jepa_loss + SIGREG_LAM * reg_loss
+        h_ctx  = tokens[:, c, :]               # context embedding
+        h_tgt  = tokens[:, c + dt, :]          # target embedding (joint training)
+        h_hat  = pred(h_ctx, torch.full((bs,), float(dt), device=dev))
+        loss   = vicreg_loss(h_hat, h_tgt, alpha=ALPHA)
         opt.zero_grad(); loss.backward(); opt.step()
 
     enc.eval()
     with torch.no_grad():
         def embed(X_np):
-            t = torch.tensor(X_np, device=dev)
+            chunks = []
-            return enc(t).mean(1).cpu().numpy()   # pool over time patches
+            for i in range(0, len(X_np), bs):
+                t = torch.tensor(X_np[i:i+bs], device=dev)
+                chunks.append(enc(t)[:, -1, :].cpu().numpy())
+            return np.concatenate(chunks, axis=0)
 
         Etr = embed(Xtr)
         Ete = embed(Xte)
 
-    # ridge linear probe (closed form)
+    # Ridge probe: fit on train, evaluate on OOS (true OOS R²)
-    A  = np.hstack([Etr, np.ones((len(Etr), 1))])
+    mu_e = Etr.mean(0); sd_e = Etr.std(0) + 1e-8
+    Etr_n = (Etr - mu_e) / sd_e
+    Ete_n = (Ete - mu_e) / sd_e
+    A  = np.hstack([Etr_n, np.ones((len(Etr_n), 1))])
     w  = np.linalg.solve(A.T @ A + 1e-3 * np.eye(A.shape[1]), A.T @ ytr)
-    pred_np = np.hstack([Ete, np.ones((len(Ete), 1))]) @ w
+    pred_np = np.hstack([Ete_n, np.ones((len(Ete_n), 1))]) @ w
     ss_res  = ((yte - pred_np) ** 2).sum()
     ss_tot  = ((yte - yte.mean()) ** 2).sum()
     val_vol_r2 = float(1 - ss_res / ss_tot)
 
+    # Phase-1: MLP supervised head on frozen embeddings
+    # Standardise targets so the head trains on unit-scale signals.
+    ytr_mu = float(ytr.mean()); ytr_sd = float(ytr.std()) + 1e-8
+    ytr_z  = (ytr - ytr_mu) / ytr_sd
+    head     = SupervisedHead(D_MODEL).to(dev)
+    head_opt = torch.optim.Adam(head.parameters(), lr=PHASE1_LR, weight_decay=1e-4)
+    Etr_t = torch.tensor(Etr_n, device=dev)
+    ytr_t = torch.tensor(ytr_z, device=dev)
+    Ete_t = torch.tensor(Ete_n, device=dev)
+    p1_bs = min(BATCH_SIZE, len(Etr_t))
+    N_tr_h = len(Etr_t)
+    # Real epoch iteration: shuffle full dataset each epoch
+    for _ in range(PHASE1_EPOCHS):
+        perm = torch.randperm(N_tr_h, device=dev)
+        for start in range(0, N_tr_h, p1_bs):
+            idx_h = perm[start:start + p1_bs]
+            loss_h = F.mse_loss(head(Etr_t[idx_h]), ytr_t[idx_h])
+            head_opt.zero_grad(); loss_h.backward(); head_opt.step()
+    head.eval()
+    with torch.no_grad():
+        pred_h_z = head(Ete_t).cpu().numpy()
+    pred_h    = pred_h_z * ytr_sd + ytr_mu          # de-standardise
+    phase1_r2 = float(1 - ((yte - pred_h) ** 2).sum() / ss_tot)
+    print("phase1_r2   = %.4f  (n_test=%d)" % (phase1_r2, len(yte)))
+
     json.dump({
-        "val_vol_r2": val_vol_r2, "n_test": len(yte),
+        "val_vol_r2": val_vol_r2, "phase1_r2": phase1_r2, "n_test": len(yte),
-        "knobs": {"WINDOW": WINDOW, "PATCH_LEN": PATCH_LEN, "STRIDE": STRIDE,
+        "knobs": {"WINDOW": WINDOW, "PATCH_LEN": PATCH_LEN,
-                  "D_MODEL": D_MODEL, "DEPTH": DEPTH, "MASK_FRAC": MASK_FRAC,
+                  "D_MODEL": D_MODEL, "DEPTH": DEPTH, "ALPHA": ALPHA,
-                  "SIGREG_LAM": SIGREG_LAM, "EPOCHS": EPOCHS},
+                  "DELTA_T_MAX": DELTA_T_MAX, "EPOCHS": EPOCHS},
     }, open("metrics.json", "w"), indent=2)
     print("val_vol_r2 = %.4f  (n_test=%d, dev=%s)" % (val_vol_r2, len(yte), dev))
 
+    # ── EXPORT BLOCK — do NOT edit (agent boundary) ──────────────────────────
+    # Set EXPORT_EMBEDDINGS=1 to write embeddings.json for the Go eval harness.
+    import os
+    if os.environ.get("EXPORT_EMBEDDINGS") == "1":
+        hourly_path2 = "data/processed/eurusd_hourly.parquet"
+        daily_path2  = "data/processed/eurusd_daily.parquet"
+        if USE_HOURLY and os.path.exists(hourly_path2):
+            df2 = pd.read_parquet(hourly_path2).reset_index(drop=True)
+            df2["date"] = pd.to_datetime(df2["datetime"])
+        else:
+            df2 = pd.read_parquet(daily_path2).reset_index(drop=True)
+            df2["date"] = pd.to_datetime(df2["date"])
+        tr_mask = df2["date"].dt.year <= 2021
+        feats2 = df2[["ret", "realized_vol"]].to_numpy(np.float32)
+        mu2 = feats2[tr_mask].mean(0); sd2 = feats2[tr_mask].std(0) + 1e-8
+        fn2 = (feats2 - mu2) / sd2
+        def _export_windows(year_mask):
+            idx = df2.index[year_mask].tolist()
+            Xs, dates, rvs = [], [], []
+            for t in idx:
+                if t - WINDOW >= 0 and t + 1 < len(df2):
+                    Xs.append(fn2[t - WINDOW:t])
+                    dates.append(str(df2["date"].iloc[t].date()))
+                    rvs.append(float(df2["realized_vol"].iloc[t + 1]))
+            if not Xs:
+                return [], [], []
+            Xa = np.stack(Xs)
+            chunks = []
+            with torch.no_grad():
+                for i in range(0, len(Xa), bs):
+                    chunks.append(enc(torch.tensor(Xa[i:i+bs], device=dev))[:, -1, :].cpu().numpy())
+            E = np.concatenate(chunks, axis=0).tolist()
+            return E, dates, rvs
+        Etr2, dates_tr, rv_tr = _export_windows(tr_mask)
+        Eoos, dates_oos, rv_oos = _export_windows(df2["date"].dt.year >= 2022)
+        hv_thr = float(np.percentile(rv_oos, 67))
+        hv_label = [1 if v >= hv_thr else 0 for v in rv_oos]
+        json.dump({"embeddings": Eoos, "dates": dates_oos,
+                   "realized_vol": rv_oos, "hv_label": hv_label,
+                   "train_embeddings": Etr2, "train_realized_vol": rv_tr},
+                  open("embeddings.json", "w"))
+        print("exported embeddings.json  train=%d  oos=%d  HV=%d/%d" % (
+            len(Etr2), len(Eoos), sum(hv_label), len(hv_label)))
+    # ── END EXPORT BLOCK ─────────────────────────────────────────────────────
+
 
 if __name__ == "__main__":
     main()