1"""Glue: grids for a symbol set, chronological session splits, origin indices, baseline fitting and scoring tables."""2from __future__ import annotations34import json5import time6from pathlib import Path78import numpy as np9import pandas as pd1011from .data.windows import SymbolGrid, WindowConfig, OriginIndex, make_example, trailing_sigma12from .eval import scores as S13from .eval.calibration import pit_histogram, reliability14from .eval.statistics import diebold_mariano, cluster_bootstrap_mean151617# ------------------------------------------------------------------------------------------------ grids & splits18def grid_dirs(cache: Path, start: str, end: str, symbols: list[str]) -> list[Path]:19 return [d for d in (cache / "grids" / f"{start}_{end}" / s for s in symbols) if d.exists()]202122def load_grids(cache: Path, start: str, end: str, symbols: list[str], mmap: bool = True) -> list[SymbolGrid]:23 return [SymbolGrid.load(d, mmap=mmap) for d in grid_dirs(cache, start, end, symbols)]242526def session_split(g: SymbolGrid, bounds: dict[str, str], warmup: int = 21) -> dict[str, np.ndarray]:27 """Session index sets per split from date bounds {train_end, val_end, test_end[, final_end]}. The first `warmup`28 sessions are never origins; each split uses sessions strictly after the previous bound."""29 dates = np.array(g.meta["session_dates"], dtype="datetime64[D]")30 idx = np.arange(len(dates))31 prev = np.datetime64("1900-01-01")32 out = {}33 for name in ("train", "val", "test", "final"):34 key = f"{name}_end"35 if key not in bounds:36 continue37 end = np.datetime64(bounds[key])38 sel = idx[(dates > prev) & (dates <= end) & (idx >= warmup)]39 out[name] = sel40 prev = end41 return out424344def build_indices(dirs: list[Path], bounds: dict[str, str], wcfg: WindowConfig, warmup: int, train_per_session: int,45 eval_stride: int, seed: int, embargo_sessions: int = 1) -> dict[str, OriginIndex]:46 """Train origins: random `train_per_session` per session; eval origins: fixed stride. An embargo of47 `embargo_sessions` sessions is removed at the start of val/test/final so that no training target overlaps."""48 grids = [SymbolGrid.load(d) for d in dirs]49 splits = [session_split(g, bounds, warmup) for g in grids]50 out = {}51 for name in splits[0].keys():52 sets = []53 for sp in splits:54 s = sp[name]55 if name != "train" and len(s) > embargo_sessions:56 s = s[embargo_sessions:]57 sets.append(s)58 if name == "train":59 out[name] = OriginIndex(dirs, wcfg, sets, stride=1, max_per_session=train_per_session, seed=seed, grids=grids)60 else:61 out[name] = OriginIndex(dirs, wcfg, sets, stride=eval_stride, grids=grids)62 return out636465# ------------------------------------------------------------------------------------------------ baselines66def origin_table(index: OriginIndex, horizons: list[int], garch_params: dict | None = None, workers: int = 8) -> pd.DataFrame:67 """Realised targets, sigma, metadata and baseline features for every origin (single causal pass, parallel)."""68 from .baselines.features import origin_frame69 return origin_frame(index, horizons, garch_params, workers)707172def samples_from_params(p: dict, n: int, seed: int = 0) -> np.ndarray:73 """Student-t / Gaussian samples [n, N] from arrays mu, scale, nu."""74 rng = np.random.default_rng(seed)75 mu, sc, nu = (np.asarray(p[k], dtype=np.float64) for k in ("mu", "scale", "nu"))76 N = len(mu)77 z = rng.standard_normal((n, N))78 nu_b = np.broadcast_to(nu, (N,))79 fin = np.isfinite(nu_b)80 g = np.ones((n, N))81 if fin.any():82 g[:, fin] = rng.chisquare(nu_b[fin], (n, fin.sum())) / nu_b[fin]83 return mu[None] + sc[None] * z / np.sqrt(g)848586# ------------------------------------------------------------------------------------------------ scoring87LEVELS = np.array([0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975])888990def score_horizon(y: np.ndarray, samples: np.ndarray, rv: np.ndarray, rv_hat: np.ndarray, sigma: np.ndarray,91 mu_hat: np.ndarray | None = None, tail_k: float = 2.5, H: int = 60,92 params: dict | None = None) -> dict[str, np.ndarray]:93 """Per-origin scores at one horizon from cumulative-return samples [n, N]."""94 out = {"crps": S.crps_samples(y, samples), "wis": S.weighted_interval_score(y, samples),95 "qloss": S.quantile_loss(y, S.quantiles_from_samples(samples, LEVELS), LEVELS),96 "mae": np.abs(y - np.median(samples, 0)), "mse": (y - samples.mean(0)) ** 2,97 "qlike": S.qlike(rv, rv_hat), "pit": S.pit_samples(y, samples), "sharpness": np.quantile(samples, 0.95, 0) - np.quantile(samples, 0.05, 0)}98 thr = tail_k * sigma * np.sqrt(H)99 event = np.abs(y) > thr100 p_tail = (np.abs(samples) > thr[None]).mean(0)101 out["brier_tail"] = S.brier(event, p_tail); out["tail_event"] = event.astype(float)102 mu = samples.mean(0) if mu_hat is None else mu_hat103 out["dir_hit"] = ((np.sign(mu) == np.sign(y)) & (y != 0)).astype(float)104 out["crps_scaled"] = out["crps"] / (sigma * np.sqrt(H))105 if params is not None and np.all(np.isfinite(params["nu"])):106 out["logscore"] = S.log_score_student_t(y, params["mu"], params["scale"], params["nu"])107 out["crps_closed"] = S.crps_student_t(y, params["mu"], params["scale"], params["nu"])108 return out109110111def summarise(per: dict[str, np.ndarray], meta: pd.DataFrame, H: int) -> dict:112 """Aggregate per-origin scores with a cluster bootstrap by trading day."""113 cl = pd.factorize(meta["session_date"])[0]114 out = {"H": H, "n": int(len(meta)), "n_days": int(len(np.unique(cl)))}115 for k in ("crps", "crps_scaled", "wis", "qloss", "mae", "mse", "qlike", "brier_tail", "sharpness", "logscore", "crps_closed"):116 if k in per:117 cb = cluster_bootstrap_mean(per[k], cl, n_boot=500)118 out[k] = cb["mean"]; out[f"{k}_lo"] = cb["lo"]; out[f"{k}_hi"] = cb["hi"]119 out["dir_acc"] = float(np.nanmean(per["dir_hit"][meta[f"cum_{H}"].values != 0]))120 out["tail_rate"] = float(per["tail_event"].mean())121 for lv in (0.5, 0.8, 0.95):122 a = 1 - lv123 pit = per["pit"]124 out[f"cov_{int(lv * 100)}"] = float(((pit >= a / 2) & (pit <= 1 - a / 2)).mean())125 ph = pit_histogram(per["pit"])126 out["pit_ks"] = ph["ks_stat"]; out["pit_chi2_p"] = ph["chi2_p"]; out["pit_var"] = ph["var"]127 return out128129130def compare(per_a: np.ndarray, per_b: np.ndarray, meta: pd.DataFrame, H: int, eval_stride: int) -> dict:131 """Model a vs b on a loss: DM with HAC lag = ceil(H/stride) overlap and day-cluster bootstrap CI of the mean diff."""132 order = pd.to_datetime(meta["session_date"]).values.astype("datetime64[m]").astype(np.int64) * 1000 + meta["m"].values133 # chronological order across symbols: sort by (date, slot); symbols interleave within a slot -> cluster by day134 lag = int(np.ceil(H / max(1, eval_stride))) * max(1, meta["symbol"].nunique())135 dm = diebold_mariano(per_a, per_b, order, lag=min(lag, len(per_a) // 4))136 cl = pd.factorize(meta["session_date"])[0]137 cb = cluster_bootstrap_mean(per_a - per_b, cl, n_boot=1000)138 return {"H": H, "diff": cb["mean"], "lo": cb["lo"], "hi": cb["hi"], "dm_t": dm["t"], "dm_p": dm["p"], "lag": dm["lag"],139 "n": int(len(per_a)), "rel_diff": float(cb["mean"] / max(1e-12, per_b.mean()))}140