"""Indexed window sampling over the dense grids: origins, three context branches and masked targets. Origin (symbol i, session s, slot m): decision time = open of slot m+1, i.e. bars 0..m of session s are complete. Eligibility: m >= m_min, m + H_max <= last target slot of the session (389, or 209 on an early close), bar m present, >= min_future_present of the H_max future slots present, s >= 1. Branch A: last L_A (390) one-minute slots ending at slot m (crosses session boundaries, flagged). Branch B: last L_B (390) completed 5-minute buckets (bucket j of session s is complete iff 5j+4 <= m). Branch C: last L_C (260) completed 30-minute buckets (complete iff 30j+29 <= m). Targets: y[h] = lc[m+h] - lc[prev present slot <= m+h-1] for h = 1..H_max with mask = present(m+h) (sum over present h <= H equals the exact H-minute log return of the last observed bar). All inputs are scaled by the trailing sigma (root mean square of present same-session 1-minute returns in the Branch-A window, floored), computed from the past only. Grids are memory-mapped and opened lazily in each DataLoader worker (the index is pickle-light).""" from __future__ import annotations import json from dataclasses import dataclass from pathlib import Path import numpy as np from .features import N_DAY, N_B5, N_B30, load_grid H_MAX = 60 SIGMA_FLOOR = 1e-5 BAR_FEATURES = ["ret", "range", "upper", "lower", "body", "close_loc", "lv", "relvol", "missing", "boundary", "gap", "minute_frac", "time_to_close", "weekday", "early_close"] F_BAR = len(BAR_FEATURES) GLOBAL_FEATURES = ["log_sigma", "overnight_gap_scaled", "minute_frac", "time_to_close", "weekday", "early_close", "rv_30_rel", "rv_day_rel", "relvol_30"] F_GLOBAL = len(GLOBAL_FEATURES) @dataclass class WindowConfig: L_A: int = 390 L_B: int = 390 L_C: int = 260 H: int = H_MAX m_min: int = 30 min_future_present: float = 0.8 extended: bool = False class SymbolGrid: """Memory-mapped per-instrument grid with light derived metadata.""" def __init__(self, symbol: str, arrays: dict, meta: dict): self.symbol, self.arrays, self.meta = symbol, arrays, meta base = arrays["base"]; mask = arrays["mask"] self.S = base.shape[0] self.n_slots = np.asarray(meta["n_slots"]) self.early = np.asarray(meta["early_close"], dtype=bool) wd = np.array([np.datetime64(d).astype("datetime64[D]").astype(int) for d in meta["session_dates"]]) self.weekday = (wd + 3) % 7 # 1970-01-01 was a Thursday -> Monday = 0 self.gap = np.array([np.nan if g is None else g for g in meta["overnight_gap"]], dtype=np.float32) self.last_target_slot = np.where(self.early, 209, 389) # flat views (reshape of a memmap is a view, no copy) self.flat_base = base[:, :N_DAY, :].reshape(-1, 5) self.flat_mask = mask[:, :N_DAY].reshape(-1) self.flat_seas = arrays["seas"][:, :N_DAY].reshape(-1) self.flat_b5 = arrays["b5"].reshape(-1, 5); self.flat_m5 = arrays["m5"].reshape(-1) self.flat_b30 = arrays["b30"].reshape(-1, 5); self.flat_m30 = arrays["m30"].reshape(-1) self.flat_lc = arrays["flat_lc"]; self.prev_present = arrays["prev_present"] self.flat_seas_r2 = arrays["seas_r2"].reshape(-1) @classmethod def load(cls, dir_: Path, mmap: bool = True) -> "SymbolGrid": arrays, meta = load_grid(Path(dir_), mmap=mmap) return cls(meta["symbol"], arrays, meta) def eligible_origins(g: SymbolGrid, cfg: WindowConfig, sessions: np.ndarray | None = None, stride: int = 1) -> np.ndarray: """(s, m) origins satisfying the eligibility rules (vectorised).""" mask = np.asarray(g.arrays["mask"][:, :N_DAY]) > 0 S = g.S sess = np.arange(1, S) if sessions is None else np.asarray([s for s in np.asarray(sessions) if 1 <= s < S], dtype=np.int64) if len(sess) == 0: return np.zeros((0, 2), dtype=np.int64) cs = np.concatenate([np.zeros((S, 1)), np.cumsum(mask, axis=1)], axis=1) # cs[s, k] = present in slots < k out = [] for s in sess: last = int(g.last_target_slot[s]) m = np.arange(cfg.m_min, last - cfg.H + 1, stride) if len(m) == 0: continue fut = (cs[s, m + 1 + cfg.H] - cs[s, m + 1]) / cfg.H ok = mask[s, m] & (fut >= cfg.min_future_present) if ok.any(): out.append(np.column_stack([np.full(ok.sum(), s), m[ok]])) return np.concatenate(out).astype(np.int64) if out else np.zeros((0, 2), dtype=np.int64) def _bar_block(b: np.ndarray, m: np.ndarray, seas: np.ndarray | None, flat_idx: np.ndarray, per_day: int, sigma: float, scale: float, g: SymbolGrid) -> np.ndarray: L = len(m) f = np.zeros((L, F_BAR), dtype=np.float32) lc = np.where(m, b[:, 3], np.nan) # previous present close within the window (forward carry of the last observed close) idx = np.where(m, np.arange(L), -1) pp = np.maximum.accumulate(idx) prev = np.full(L, np.nan) prev[1:] = np.where(pp[:-1] >= 0, lc[np.clip(pp[:-1], 0, None)], np.nan) ret = np.where(m & ~np.isnan(prev), lc - prev, 0.0) s_eff = sigma * scale f[:, 0] = ret / s_eff with np.errstate(invalid="ignore"): rng = np.where(m, b[:, 1] - b[:, 2], 0.0) f[:, 1] = rng / s_eff f[:, 2] = np.where(m, b[:, 1] - np.maximum(b[:, 0], b[:, 3]), 0.0) / s_eff f[:, 3] = np.where(m, np.minimum(b[:, 0], b[:, 3]) - b[:, 2], 0.0) / s_eff f[:, 4] = np.where(m, b[:, 3] - b[:, 0], 0.0) / s_eff f[:, 5] = np.where(m & (rng > 0), (b[:, 3] - b[:, 2]) / np.where(rng > 0, rng, 1.0), 0.5) - 0.5 f[:, 6] = np.where(m, b[:, 4], 0.0) / 10.0 f[:, 7] = np.where(m & ~np.isnan(seas), b[:, 4] - seas, 0.0) if seas is not None else 0.0 f[:, 8] = (~m).astype(np.float32) sess = flat_idx // per_day slot = flat_idx % per_day boundary = np.zeros(L, dtype=np.float32) boundary[1:] = (sess[1:] != sess[:-1]).astype(np.float32) f[:, 9] = boundary width = 390 // per_day gap_min = np.full(L, float(width)) if L > 1: cross = sess[1:] != sess[:-1] # minutes from the previous row's bar open to this row's bar open across the session change rem_prev = (g.n_slots[sess[:-1]] - 1) - (slot[:-1] + 1) * width + 1 gap_cross = width + rem_prev + (24 * 60 - 390) * np.maximum(1, sess[1:] - sess[:-1]) gap_min[1:] = np.where(cross, gap_cross, width) f[:, 10] = np.log1p(gap_min) / 8.0 f[:, 11] = (slot * width) / 389.0 f[:, 12] = (g.n_slots[sess] - 1 - slot * width) / 389.0 f[:, 13] = g.weekday[sess] / 4.0 f[:, 14] = g.early[sess].astype(np.float32) return np.nan_to_num(f, nan=0.0, posinf=0.0, neginf=0.0) def trailing_returns(g: SymbolGrid, s: int, m: int, L: int = 390) -> np.ndarray: """Same-session 1-minute returns between present bars inside the last L+1 slots ending at (s, m).""" end = s * N_DAY + m + 1 lo = max(0, end - L - 1) idx = np.arange(lo, end) lc = np.asarray(g.flat_lc[idx], dtype=np.float64) ok = ~np.isnan(lc) if ok.sum() < 2: return np.zeros(0) r = np.diff(lc[ok]); sess = idx[ok] // N_DAY return r[sess[1:] == sess[:-1]] def trailing_sigma(g: SymbolGrid, s: int, m: int, L: int = 390) -> tuple[float, np.ndarray]: r = trailing_returns(g, s, m, L) sigma = float(np.sqrt(np.mean(r ** 2))) if len(r) else SIGMA_FLOOR return max(sigma, SIGMA_FLOOR), r def make_example(g: SymbolGrid, s: int, m: int, cfg: WindowConfig) -> dict[str, np.ndarray]: sigma, r_hist = trailing_sigma(g, s, m, cfg.L_A) endA = s * N_DAY + m + 1 idxA = np.arange(endA - cfg.L_A, endA); validA = idxA >= 0; idxA_c = np.clip(idxA, 0, None) bA = np.asarray(g.flat_base[idxA_c]); mA = (np.asarray(g.flat_mask[idxA_c]) > 0) & validA xA = _bar_block(bA, mA, np.asarray(g.flat_seas[idxA_c]), idxA_c, N_DAY, sigma, 1.0, g) xA[~validA] = 0.0; xA[~validA, 8] = 1.0 jB = (m + 1) // 5 - 1 endB = s * N_B5 + jB + 1 idxB = np.arange(endB - cfg.L_B, endB); validB = idxB >= 0; idxB_c = np.clip(idxB, 0, None) bB = np.asarray(g.flat_b5[idxB_c]); mB = (np.asarray(g.flat_m5[idxB_c]) > 0) & validB xB = _bar_block(bB, mB, None, idxB_c, N_B5, sigma, np.sqrt(5.0), g) xB[~validB] = 0.0; xB[~validB, 8] = 1.0 jC = (m + 1) // 30 - 1 endC = s * N_B30 + jC + 1 idxC = np.arange(endC - cfg.L_C, endC); validC = idxC >= 0; idxC_c = np.clip(idxC, 0, None) bC = np.asarray(g.flat_b30[idxC_c]); mC = (np.asarray(g.flat_m30[idxC_c]) > 0) & validC xC = _bar_block(bC, mC, None, idxC_c, N_B30, sigma, np.sqrt(30.0), g) xC[~validC] = 0.0; xC[~validC, 8] = 1.0 base_idx = s * N_DAY + m tgt_idx = base_idx + np.arange(1, cfg.H + 1) lc_t = np.asarray(g.flat_lc[tgt_idx], dtype=np.float64) prev_idx = np.asarray(g.prev_present[tgt_idx - 1]) y_mask = ~np.isnan(lc_t) & (prev_idx >= 0) y = np.where(y_mask, lc_t - np.asarray(g.flat_lc[np.clip(prev_idx, 0, None)], dtype=np.float64), 0.0).astype(np.float32) fut_slot = m + np.arange(1, cfg.H + 1) fut_cal = np.stack([fut_slot / 389.0, (g.n_slots[s] - 1 - fut_slot) / 389.0, np.full(cfg.H, g.weekday[s] / 4.0), np.full(cfg.H, float(g.early[s]))], axis=1).astype(np.float32) rv30 = float(np.mean(r_hist[-30:] ** 2)) if len(r_hist) >= 5 else sigma ** 2 rvday = float(np.mean(r_hist[-m:] ** 2)) if m >= 5 and len(r_hist) >= 5 else sigma ** 2 relvol30 = float(np.mean(xA[-30:, 7])) gap = g.gap[s] glob = np.array([np.log(sigma), 0.0 if np.isnan(gap) else gap / (sigma * 20), m / 389.0, (g.n_slots[s] - 1 - m) / 389.0, g.weekday[s] / 4.0, float(g.early[s]), np.log(rv30 / sigma ** 2 + 1e-12), np.log(rvday / sigma ** 2 + 1e-12), relvol30], dtype=np.float32) return {"xA": xA, "mA": mA.astype(np.uint8), "xB": xB, "mB": mB.astype(np.uint8), "xC": xC, "mC": mC.astype(np.uint8), "y": y, "y_mask": y_mask.astype(np.uint8), "fut_cal": fut_cal, "glob": np.nan_to_num(glob), "sigma": np.float32(sigma), "origin": np.array([s, m], dtype=np.int64)} class OriginIndex: """Flat index of (symbol, session, slot) origins across many grids. Holds grid directories; grids are opened lazily (memory-mapped) so the index can be pickled to DataLoader workers cheaply.""" def __init__(self, grid_dirs: list[Path], cfg: WindowConfig, session_sets: list[np.ndarray] | None = None, stride: int = 1, max_per_session: int | None = None, seed: int = 0, grids: list[SymbolGrid] | None = None): self.grid_dirs = [Path(d) for d in grid_dirs] self.cfg = cfg self._grids: dict[int, SymbolGrid] = {} if grids is not None: self._grids = dict(enumerate(grids)) rng = np.random.default_rng(seed) rows = [] for i in range(len(self.grid_dirs)): g = self.grid(i) o = eligible_origins(g, cfg, None if session_sets is None else session_sets[i], stride) if max_per_session is not None and len(o): # vectorised per-session subsampling: random key, sort by (session, key), keep the first k per session keys = rng.random(len(o)) order = np.lexsort((keys, o[:, 0])) os_ = o[order] starts = np.r_[0, np.flatnonzero(np.diff(os_[:, 0])) + 1] sizes = np.diff(np.r_[starts, len(os_)]) rank = np.arange(len(os_)) - np.repeat(starts, sizes) o = os_[rank < max_per_session] o = o[np.lexsort((o[:, 1], o[:, 0]))] if len(o): rows.append(np.column_stack([np.full(len(o), i), o])) self.rows = np.concatenate(rows).astype(np.int64) if rows else np.zeros((0, 3), dtype=np.int64) self.symbols = [g.symbol for g in (self.grid(i) for i in range(len(self.grid_dirs)))] def grid(self, i: int) -> SymbolGrid: g = self._grids.get(i) if g is None: g = SymbolGrid.load(self.grid_dirs[i], mmap=True) self._grids[i] = g return g @property def grids(self) -> list[SymbolGrid]: return [self.grid(i) for i in range(len(self.grid_dirs))] def __getstate__(self): d = self.__dict__.copy(); d["_grids"] = {} return d def __len__(self): return len(self.rows) def example(self, k: int) -> dict[str, np.ndarray]: i, s, m = self.rows[k] ex = make_example(self.grid(int(i)), int(s), int(m), self.cfg) ex["sym"] = np.int64(i) return ex def save_index(idx: OriginIndex, path: Path) -> None: np.save(path, idx.rows)