Files
jepa-fx-risk/train.py
T
mathiasandClaude Sonnet 4.6 f5ce8d6706
CD / Lint / Test / Vet (push) Failing after 4s
CD / Build & Import (push) Has been skipped
CD / Deploy via GitOps (push) Has been skipped
chore(loop): 6 more iters — plateau at ~0.34-0.37 (1/6 kept)
Toy encoder near ceiling. 1 kept (val_vol_r2 0.3032→0.3442), 5 reverts.
Consistent plateau = time to swap in TS-JEPA backbone (#3/#5).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-24 07:37:29 +02:00

93 lines
3.2 KiB
Python

"""train.py — the ONLY file the autoresearch agent may edit (Phase-1 contract).
Toy slice: a tiny self-supervised encoder (masked reconstruction of windowed
daily [return, realized_vol]) → FROZEN → linear probe predicts NEXT-day realized
vol → val_vol_r2 = OOS R². The agent improves val_vol_r2 by editing the encoder /
objective / masking below. Writes metrics.json (the scalar the loop reads).
python train.py
"""
import json
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
# --- agent-tunable knobs ---
WINDOW = 20
EMBED_DIM = 64
MASK_FRAC = 0.40
EPOCHS = 200
LR = 1e-3
SEED = 0
# ---------------------------
torch.manual_seed(SEED)
np.random.seed(SEED)
dev = "cuda" if torch.cuda.is_available() else "cpu"
def build():
df = pd.read_parquet("data/processed/eurusd_daily.parquet").reset_index(drop=True)
feats = df[["ret", "realized_vol"]].to_numpy(np.float32)
target = df["realized_vol"].to_numpy(np.float32) # predict NEXT-day RV
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)) # time-ordered OOS split
mu, sd = X[:n_tr].mean((0, 1)), X[:n_tr].std((0, 1)) + 1e-8 # train-only stats
X = (X - mu) / sd
return (X[:n_tr], y[:n_tr]), (X[n_tr:], y[n_tr:])
class Encoder(nn.Module):
def __init__(self, win, emb):
super().__init__()
self.net = nn.Sequential(
nn.Flatten(),
nn.Linear(win * 2, 256),
nn.LayerNorm(256),
nn.GELU(),
nn.Linear(256, emb)
)
def forward(self, x):
return self.net(x)
def main():
(Xtr, ytr), (Xte, yte) = build()
Xtr_t = torch.tensor(Xtr, device=dev)
enc = Encoder(WINDOW, EMBED_DIM).to(dev)
dec = nn.Sequential(nn.Linear(EMBED_DIM, 256), nn.GELU(), nn.Linear(256, WINDOW * 2)).to(dev)
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=LR)
for _ in range(EPOCHS): # SSL: masked reconstruction of the window
mask = (torch.rand_like(Xtr_t) > MASK_FRAC).float()
rec = dec(enc((Xtr_t * mask)))
loss = (((rec - Xtr_t.flatten(1)) ** 2) * (1 - mask.flatten(1))).mean()
opt.zero_grad(); loss.backward(); opt.step()
enc.eval()
with torch.no_grad(): # FROZEN embeddings
Etr = enc(Xtr_t).cpu().numpy()
Ete = enc(torch.tensor(Xte, device=dev)).cpu().numpy()
# linear probe (ridge, closed form) on frozen embeddings → val_vol_r2 (OOS R²)
A = np.hstack([Etr, np.ones((len(Etr), 1))])
w = np.linalg.solve(A.T @ A + 1e-3 * np.eye(A.shape[1]), A.T @ ytr)
pred = np.hstack([Ete, np.ones((len(Ete), 1))]) @ w
ss_res = ((yte - pred) ** 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, "EMBED_DIM": EMBED_DIM, "MASK_FRAC": MASK_FRAC, "EPOCHS": EPOCHS}},
open("metrics.json", "w"), indent=2)
print("val_vol_r2 = %.4f (n_test=%d, dev=%s)" % (val_vol_r2, len(yte), dev))
if __name__ == "__main__":
main()