generated from mathias/template-go-web
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
  mathiasandClaude Sonnet 4.6
|
bde651b0df | ||
|
mathiasandClaude Sonnet 4.6
|
20aeecb971 |
+84
-34
@@ -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", ...]}
|
||||
// Generated by train.py when run with EXPORT_EMBEDDINGS=1.
|
||||
// embeddings.json format (from train.py 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 {
|
||||
log.Error("silhouette requires HV labels in embeddings.json")
|
||||
if len(d.HVLabel) == 0 {
|
||||
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"`
|
||||
Dates []string `json:"dates"`
|
||||
RealizedVol []float64 `json:"realized_vol"`
|
||||
HVLabel []int `json:"hv_label"`
|
||||
Embeddings [][]float64 `json:"embeddings"`
|
||||
Dates []string `json:"dates"`
|
||||
RealizedVol []float64 `json:"realized_vol"`
|
||||
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
|
||||
n, dim := len(d.Embeddings), len(d.Embeddings[0])
|
||||
// standardise centres + scales to zero mean / unit std; returns normalised rows.
|
||||
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 j, v := range row {
|
||||
for _, r := range rows {
|
||||
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 j, v := range row {
|
||||
for _, r := range rows {
|
||||
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)
|
||||
for i, row := range d.Embeddings {
|
||||
norm[i] = make([]float64, dim)
|
||||
for j, v := range row {
|
||||
norm[i][j] = (v - mu[j]) / sd[j]
|
||||
out := make([][]float64, n)
|
||||
for i, r := range rows {
|
||||
out[i] = make([]float64, dim)
|
||||
for j, v := range r {
|
||||
out[i][j] = (v - mu[j]) / sd[j]
|
||||
}
|
||||
}
|
||||
return out, mu, sd
|
||||
}
|
||||
|
||||
if len(d.HVLabel) > 0 {
|
||||
labels = d.HVLabel
|
||||
// applyStandardise normalises rows using pre-computed mu and sd.
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Failing tests for HEPA backbone 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 new backbone must satisfy BEFORE implementation.
|
||||
"""
|
||||
import math
|
||||
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():
|
||||
import importlib.util, sys
|
||||
spec = importlib.util.spec_from_file_location("train", "train.py")
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
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() still returns year-based OOS split (2022-2023)
|
||||
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 set should be ~600 windows (2 years of daily data)
|
||||
assert 400 < len(Xte) < 900, f"OOS size unexpected: {len(Xte)}"
|
||||
@@ -1,14 +1,13 @@
|
||||
"""train.py — autoresearch agent file (only this may be edited).
|
||||
|
||||
TS-JEPA backbone with SIGReg regularization (Balestriero & LeCun, LeJEPA
|
||||
arXiv:2511.08544; time-series placement from ChronoJEPA arXiv: 2505.XXXXX).
|
||||
HEPA backbone (Petersen et al., arXiv:2605.11130, ICML 2026 Spotlight):
|
||||
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 →
|
||||
linear probe predicts NEXT-day realized vol → val_vol_r2 (OOS R²).
|
||||
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.
|
||||
Agent may tune: encoder depth/width, patch geometry, ALPHA, DELTA_T_MAX,
|
||||
optimizer, LR. Do NOT touch prepare_data.py, loop.py, or the data pipeline.
|
||||
"""
|
||||
import json
|
||||
import math
|
||||
@@ -16,19 +15,19 @@ import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# --- agent-tunable knobs ---
|
||||
WINDOW = 60 # INCREASED lookback for better volatility persistence capture
|
||||
PATCH_LEN = 5 # time-patch size (must divide WINDOW)
|
||||
STRIDE = 5
|
||||
D_MODEL = 64 # transformer hidden dim - INCREASED for capacity
|
||||
DEPTH = 2 # transformer layers
|
||||
N_HEADS = 4
|
||||
MASK_FRAC = 0.50 # INCREASED mask fraction to force the encoder to learn better global representations
|
||||
SIGREG_LAM = 0.01 # SIGReg weight (λ) - REDUCED to allow more representation capacity
|
||||
EPOCHS = 300
|
||||
LR = 3e-4
|
||||
SEED = 0
|
||||
WINDOW = 60
|
||||
PATCH_LEN = 10 # non-overlapping patches (6 tokens per window)
|
||||
D_MODEL = 128
|
||||
DEPTH = 2
|
||||
N_HEADS = 4
|
||||
ALPHA = 0.1 # VICReg mixing weight (fixed at 0.1 in HEPA paper)
|
||||
DELTA_T_MAX = 3 # max prediction horizon in patches (1..min(DELTA_T_MAX, N-1-c))
|
||||
EPOCHS = 300
|
||||
LR = 3e-4
|
||||
SEED = 0
|
||||
# ---------------------------
|
||||
|
||||
torch.manual_seed(SEED)
|
||||
@@ -36,139 +35,197 @@ 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:
|
||||
"""Epps-Pulley test statistic pushes token embeddings toward isotropic Gaussian.
|
||||
def vicreg_loss(h_pred: torch.Tensor, h_target: torch.Tensor, alpha: float = 0.1) -> torch.Tensor:
|
||||
"""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
|
||||
z = tokens.reshape(B * T, D) # (N, D)
|
||||
t = torch.linspace(0, 3, knots, device=z.device, dtype=z.float().dtype)
|
||||
dt = 3.0 / (knots - 1)
|
||||
w = torch.full((knots,), 2 * dt, device=z.device, dtype=z.float().dtype)
|
||||
w[0] = dt; w[-1] = dt
|
||||
phi = torch.exp(-t.square() / 2.0)
|
||||
|
||||
A = torch.randn(D, 256, device=z.device, dtype=z.float().dtype)
|
||||
A = A / A.norm(p=2, dim=0)
|
||||
x_t = (z.float() @ A).unsqueeze(-1) * t # (N, 256, knots)
|
||||
err = (x_t.cos().mean(0) - phi).square() + x_t.sin().mean(0).square()
|
||||
return ((err @ (w * phi)) * z.shape[0]).mean()
|
||||
pred_n = F.normalize(h_pred, dim=-1)
|
||||
targ_n = F.normalize(h_target, dim=-1)
|
||||
l1 = F.l1_loss(pred_n, targ_n)
|
||||
# variance hinge: push each feature std toward ≥ 1
|
||||
std = h_pred.std(dim=0) + 1e-4
|
||||
l_var = F.relu(1.0 - std).mean()
|
||||
# covariance penalty: decorrelate features
|
||||
B, D = h_pred.shape
|
||||
h_c = h_pred - h_pred.mean(dim=0, keepdim=True)
|
||||
cov = (h_c.t() @ h_c) / max(B - 1, 1)
|
||||
off = cov - torch.diag(torch.diag(cov))
|
||||
l_cov = (off ** 2).sum() / D
|
||||
return (1 - alpha) * l1 + alpha * (l_var + l_cov)
|
||||
|
||||
|
||||
# ── Encoder + Predictor ─────────────────────────────────────────────────────
|
||||
# ── CausalEncoder ─────────────────────────────────────────────────────────────
|
||||
|
||||
class PatchEncoder(nn.Module):
|
||||
"""PatchTST-style encoder for univariate windows."""
|
||||
def __init__(self, in_feats, patch_len, stride, d_model, depth, n_heads):
|
||||
class CausalEncoder(nn.Module):
|
||||
"""Non-overlapping patches → per-patch LayerNorm → causal Transformer → all tokens (B, N, D).
|
||||
|
||||
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)
|
||||
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)
|
||||
self.tf = nn.TransformerEncoder(layer, num_layers=depth)
|
||||
self.norm = nn.LayerNorm(d_model)
|
||||
|
||||
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
|
||||
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) # (B, T, D)
|
||||
P = self.patch_len
|
||||
N = W // P
|
||||
tokens = x[:, :N * P, :].reshape(B, N, P * F)
|
||||
tokens = self.embed(self.patch_norm(tokens))
|
||||
# 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):
|
||||
def __init__(self, d_model):
|
||||
# ── HorizonPredictor ─────────────────────────────────────────────────────────
|
||||
|
||||
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(),
|
||||
nn.Linear(d_model, d_model))
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(d_model + 1, d_model), nn.GELU(),
|
||||
nn.Linear(d_model, d_model), nn.GELU(),
|
||||
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 ────────────────────────────────────────────────────────────────────
|
||||
# ── Data ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def build():
|
||||
"""Year-based split: encoder trains on 2019-2021; probe evaluates on 2022-2023 OOS."""
|
||||
df = pd.read_parquet("data/processed/eurusd_daily.parquet").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 = [], []
|
||||
for t in range(WINDOW, len(df) - 1):
|
||||
X.append(feats[t - WINDOW:t])
|
||||
y.append(target[t + 1])
|
||||
X = np.stack(X); y = np.array(y, np.float32)
|
||||
n_tr = int(0.7 * len(X))
|
||||
mu = X[:n_tr].mean((0, 1))
|
||||
sd = X[:n_tr].std((0, 1)) + 1e-8
|
||||
X = (X - mu) / sd
|
||||
return (X[:n_tr], y[:n_tr]), (X[n_tr:], y[n_tr:])
|
||||
tr_idx = df.index[df["date"].dt.year <= 2021].tolist()
|
||||
te_idx = df.index[df["date"].dt.year >= 2022].tolist()
|
||||
mu = feats[:tr_idx[-1]+1].mean(0)
|
||||
sd = feats[:tr_idx[-1]+1].std(0) + 1e-8
|
||||
fn = (feats - mu) / sd
|
||||
def windows(idx):
|
||||
X, y = [], []
|
||||
for t in idx:
|
||||
if t - WINDOW >= 0 and t + 1 < len(df):
|
||||
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]
|
||||
Xtr_t = torch.tensor(Xtr, device=dev)
|
||||
enc = PatchEncoder(n_feats, PATCH_LEN, STRIDE, D_MODEL, DEPTH, N_HEADS).to(dev)
|
||||
pred = Predictor(D_MODEL).to(dev)
|
||||
opt = torch.optim.AdamW(list(enc.parameters()) + list(pred.parameters()), lr=LR)
|
||||
n_feats = Xtr.shape[2]
|
||||
n_patches = WINDOW // PATCH_LEN
|
||||
Xtr_t = torch.tensor(Xtr, device=dev)
|
||||
|
||||
n_patches = (WINDOW - PATCH_LEN) // STRIDE + 1
|
||||
n_mask = max(1, int(MASK_FRAC * n_patches))
|
||||
enc = CausalEncoder(n_feats, PATCH_LEN, D_MODEL, N_HEADS, DEPTH).to(dev)
|
||||
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
|
||||
idx_mask = torch.randperm(n_patches)[:n_mask]
|
||||
ctx_mask = torch.ones(n_patches, dtype=torch.bool, device=dev)
|
||||
ctx_mask[idx_mask] = False
|
||||
# Sample random context position and horizon; Δt log-biased toward short
|
||||
c = torch.randint(0, n_patches - 1, ()).item()
|
||||
dt = torch.randint(1, max(2, min(DELTA_T_MAX, n_patches - 1 - c) + 1), ()).item()
|
||||
|
||||
tokens_ctx = enc(Xtr_t) # encode all (B, T, D)
|
||||
tokens_target = enc(Xtr_t).detach() # target (frozen): same input, no grad
|
||||
pred_out = pred(tokens_ctx[:, idx_mask, :])
|
||||
jepa_loss = ((pred_out - tokens_target[:, idx_mask, :]) ** 2).mean()
|
||||
reg_loss = sigreg(tokens_ctx)
|
||||
loss = jepa_loss + SIGREG_LAM * reg_loss
|
||||
tokens = enc(Xtr_t) # (B, N, D)
|
||||
h_ctx = tokens[:, c, :] # context embedding
|
||||
h_tgt = tokens[:, c + dt, :] # target embedding (joint training)
|
||||
h_hat = pred(h_ctx, torch.full((len(Xtr),), 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)
|
||||
return enc(t).mean(1).cpu().numpy() # pool over time patches
|
||||
return enc(t)[:, -1, :].cpu().numpy() # last token = full-context summary
|
||||
|
||||
Etr = embed(Xtr)
|
||||
Ete = embed(Xte)
|
||||
|
||||
# ridge linear probe (closed form)
|
||||
A = np.hstack([Etr, np.ones((len(Etr), 1))])
|
||||
# Ridge probe: fit on train, evaluate on OOS (true OOS R²)
|
||||
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)
|
||||
|
||||
json.dump({
|
||||
"val_vol_r2": val_vol_r2, "n_test": len(yte),
|
||||
"knobs": {"WINDOW": WINDOW, "PATCH_LEN": PATCH_LEN, "STRIDE": STRIDE,
|
||||
"D_MODEL": D_MODEL, "DEPTH": DEPTH, "MASK_FRAC": MASK_FRAC,
|
||||
"SIGREG_LAM": SIGREG_LAM, "EPOCHS": EPOCHS},
|
||||
"knobs": {"WINDOW": WINDOW, "PATCH_LEN": PATCH_LEN,
|
||||
"D_MODEL": D_MODEL, "DEPTH": DEPTH, "ALPHA": ALPHA,
|
||||
"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":
|
||||
df2 = pd.read_parquet("data/processed/eurusd_daily.parquet").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:
|
||||
Xs.append(fn2[t - WINDOW:t])
|
||||
dates.append(str(df2["date"].iloc[t].date()))
|
||||
rvs.append(float(df2["realized_vol"].iloc[t]))
|
||||
if not Xs:
|
||||
return [], [], []
|
||||
with torch.no_grad():
|
||||
E = enc(torch.tensor(np.stack(Xs), device=dev))[:, -1, :].cpu().numpy().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()
|
||||
Reference in New Issue
Block a user