1"""End-to-end experiment: grids -> splits -> baselines + neural variants -> predictions -> scores -> tables."""2from __future__ import annotations34import json5import time6from pathlib import Path78import numpy as np9import pandas as pd10import torch1112from .config import load_config13from .data.windows import WindowConfig, OriginIndex, make_example14from .models.model import MRPT, MRPTConfig15from .baselines.har import HARRV16from .baselines.trees import RidgeBaseline, ElasticNetBaseline, LightGBMBaseline17from .baselines.garch import five_minute_returns18from .baselines.neural import MLPBaseline, TCNBaseline19from .train.trainer import train, pick_device20from .train.sampling import OriginDataset21from .inference.predict import predict_index22from .inference.serialization import save_predictions, write_json23from . import pipeline as P2425ROOT = Path(__file__).resolve().parents[2]262728def symbols_for(man: dict, which: str, roles=("pooled",)) -> list[dict]:29 inst = {r["symbol"]: r for r in man["instruments"]}30 if which == "pilot":31 syms = man["pilot_symbols"]32 elif which == "feasibility":33 syms = man["feasibility_symbols"]34 else:35 syms = [s for s, r in inst.items() if r["role"] in roles]36 return [inst[s] for s in syms if inst[s]["role"] in roles]373839# ------------------------------------------------------------------------------------------------ baselines40def _garch_fit(r5: np.ndarray) -> dict | None:41 from arch import arch_model42 import warnings43 if len(r5) < 500:44 return None45 y = r5 * 1e446 with warnings.catch_warnings():47 warnings.simplefilter("ignore")48 res = arch_model(y, mean="Zero", vol="GARCH", p=1, q=1, dist="t").fit(disp="off", show_warning=False)49 p = res.params50 return {"omega": float(p["omega"]) / 1e8, "alpha": float(p["alpha[1]"]), "beta": float(p["beta[1]"]),51 "nu": float(np.clip(p["nu"], 2.5, 50)), "n_obs": int(len(r5))}525354def fit_garch_params(idx_train: OriginIndex) -> dict[int, dict]:55 """GARCH(1,1)-t per instrument on the 5-minute returns of its training sessions."""56 out = {}57 for i in np.unique(idx_train.rows[:, 0]):58 s_tr = idx_train.rows[idx_train.rows[:, 0] == i, 1]59 pr = _garch_fit(five_minute_returns(idx_train.grid(int(i)), int(s_tr.min()), int(s_tr.max())))60 if pr is not None:61 out[int(i)] = pr62 return out636465def fit_and_predict_baselines(tab_tr: pd.DataFrame, tab_ev: pd.DataFrame, horizons: list[int], out_dir: Path, seed: int,66 which: tuple[str, ...]) -> dict[str, dict[int, dict]]:67 """Baselines from the origin frames (training frame for fitting, evaluation frame for prediction)."""68 from .baselines.features import TREE_FEATURES, HAR_FEATURES69 t0 = time.time()70 N = len(tab_ev); sig = tab_ev["sigma"].values71 preds: dict[str, dict[int, dict]] = {}72 inf = np.full(N, np.inf); zero = np.zeros(N)73 if "naive" in which:74 preds["naive_rw"] = {H: {"mu": zero, "scale": sig * np.sqrt(H), "nu": inf, "var": sig ** 2 * H} for H in horizons}75 preds["seasonal_vol"] = {H: {"mu": zero, "scale": np.sqrt(tab_ev[f"seas_var_{H}"].values), "nu": inf,76 "var": tab_ev[f"seas_var_{H}"].values} for H in horizons}77 nu = 5.078 preds["ewma"] = {H: {"mu": zero, "scale": np.sqrt(tab_ev["ewma_var1"].values * H * (nu - 2) / nu), "nu": np.full(N, nu),79 "var": tab_ev["ewma_var1"].values * H} for H in horizons}80 if "har" in which:81 har = HARRV(); preds["har_rv"] = {}82 Xtr = tab_tr[HAR_FEATURES].values; Xev = tab_ev[HAR_FEATURES].values83 for H in horizons:84 har.fit(Xtr, tab_tr["m"].values, tab_tr[f"rv_{H}"].values, tab_tr[f"cum_{H}"].values, H)85 preds["har_rv"][H] = har.predict_batch(Xev, tab_ev["m"].values, H)86 if "garch" in which and f"garch_var_{horizons[0]}" in tab_ev.columns:87 preds["garch_t"] = {}88 for H in horizons:89 var = tab_ev[f"garch_var_{H}"].values.astype(float); nu = tab_ev["garch_nu"].values.astype(float)90 miss = ~np.isfinite(var) | (var <= 0)91 var = np.where(miss, sig ** 2 * H, var); nu = np.where(np.isfinite(nu), nu, 6.0)92 preds["garch_t"][H] = {"mu": zero, "scale": np.sqrt(var * (nu - 2) / nu), "nu": nu, "var": var}93 if "trees" in which:94 Xtr = tab_tr[TREE_FEATURES].values; Xev = tab_ev[TREE_FEATURES].values95 for name, model in (("ridge", RidgeBaseline()), ("elastic_net", ElasticNetBaseline()), ("lightgbm", LightGBMBaseline(seed=seed))):96 preds[name] = {}97 for H in horizons:98 model.fit(Xtr, tab_tr[f"cum_{H}"].values, tab_tr[f"rv_{H}"].values, tab_tr["sigma"].values, H)99 preds[name][H] = model.predict_batch(Xev, sig, H)100 write_json({"wall_s": time.time() - t0, "n_train": int(len(tab_tr)), "n_eval": N, "baselines": list(preds)}, out_dir / "baselines_meta.json")101 return preds102103104# ------------------------------------------------------------------------------------------------ neural105def build_model(spec: dict, n_symbols: int) -> torch.nn.Module:106 kind = spec.get("kind", "mrpt")107 if kind == "mlp":108 return MLPBaseline()109 if kind == "tcn":110 return TCNBaseline()111 kw = dict(spec.get("kwargs", {}))112 if kw.pop("ticker_embedding", False):113 kw["n_symbols"] = n_symbols114 return MRPT(MRPTConfig.preset(spec.get("preset", "small"), **kw))115116117def train_variant(spec: dict, idx_train: OriginIndex, idx_val: OriginIndex, cfg: dict, out_dir: Path, n_symbols: int) -> tuple[torch.nn.Module, dict]:118 tcfg = cfg["train"]119 model = build_model(spec, n_symbols)120 ds_tr = OriginDataset(idx_train, spec.get("feature_set", "full"), spec.get("shuffle_minutes", False))121 ds_va = OriginDataset(idx_val, spec.get("feature_set", "full"))122 summ = train(model, ds_tr, ds_va, out_dir, batch_size=tcfg["batch_size"], max_steps=spec.get("max_steps", tcfg["max_steps"]),123 lr=tcfg["lr"], warmup_steps=tcfg["warmup_steps"], weight_decay=tcfg["weight_decay"], grad_clip=tcfg["grad_clip"],124 eval_every=tcfg["eval_every"], patience=tcfg.get("patience", 5), seed=spec.get("seed", tcfg["seed"]),125 num_workers=tcfg["num_workers"], aux_vol_weight=tcfg.get("aux_vol_weight", 0.0),126 val_max_batches=tcfg.get("val_max_batches", 50))127 # Stage B: copula head trained after the marginals, marginals frozen128 if getattr(model, "copula", None) is not None:129 from .train import checkpoint as ck130 ck.load(out_dir / "best.pt", model)131 summ_b = train(model, ds_tr, ds_va, out_dir / "copula", batch_size=tcfg["batch_size"], max_steps=spec.get("copula_steps", 300),132 lr=tcfg["lr"], warmup_steps=20, weight_decay=0.0, grad_clip=tcfg["grad_clip"], eval_every=100, patience=3,133 seed=spec.get("seed", tcfg["seed"]), num_workers=tcfg["num_workers"], copula_only=True, resume=False)134 summ["copula"] = summ_b135 return model, summ136137138# ------------------------------------------------------------------------------------------------ scoring139def score_predictions(name: str, tab: pd.DataFrame, horizons: list[int], out_dir: Path, n_samples: int, seed: int,140 neural: dict | None = None, param_preds: dict | None = None, universe: dict | None = None) -> dict:141 """Per-origin scores at each horizon (saved) and summaries. Either `neural` (predict_index output) or142 `param_preds` {H: {mu, scale, nu, var}} must be given."""143 from .eval.subgroups import add_subgroups144 meta = add_subgroups(tab, universe) if universe else tab145 summaries, per_all = {}, {}146 for H in horizons:147 y = tab[f"cum_{H}"].values; rv = tab[f"rv_{H}"].values; sigma = tab["sigma"].values148 if neural is not None:149 samples = neural[f"cum_samples_{H}"]; rv_hat = neural[f"rv_samples_{H}"].mean(0)150 # consistency: realised from the predictor must equal the origin table151 assert np.allclose(neural[f"cum_{H}"], y, atol=1e-6), "origin/prediction misalignment"152 mu_hat = samples.mean(0); params = None153 else:154 p = param_preds[H]; samples = P.samples_from_params(p, n_samples, seed); rv_hat = p["var"]; mu_hat = p["mu"]155 params = p if np.all(np.isfinite(p["nu"])) else None156 per = P.score_horizon(y, samples, rv, rv_hat, sigma, mu_hat, H=H, params=params)157 per_all[H] = per158 summaries[H] = P.summarise(per, meta, H)159 np.savez_compressed(out_dir / f"per_origin_{name}.npz", **{f"{k}_{H}": v for H, per in per_all.items() for k, v in per.items()})160 write_json(summaries, out_dir / f"summary_{name}.json")161 return {"summary": summaries, "per": per_all}162163164def comparisons(scored: dict[str, dict], tab: pd.DataFrame, horizons: list[int], stride: int, metric: str, reference: str) -> pd.DataFrame:165 rows = []166 for name, sc in scored.items():167 if name == reference:168 continue169 for H in horizons:170 r = P.compare(sc["per"][H][metric], scored[reference]["per"][H][metric], tab, H, stride)171 rows.append({"model": name, "reference": reference, "metric": metric, **r})172 return pd.DataFrame(rows)173174175def leaderboard(scored: dict[str, dict], horizons: list[int]) -> pd.DataFrame:176 rows = []177 for name, sc in scored.items():178 for H in horizons:179 s = sc["summary"][H]180 rows.append({"model": name, "H": H, **{k: s[k] for k in s if k not in ("H",)}})181 return pd.DataFrame(rows)182