1"""Indexed window sampling over the dense grids: origins, three context branches and masked targets.23Origin (symbol i, session s, slot m): decision time = open of slot m+1, i.e. bars 0..m of session s are complete.4Eligibility: m >= m_min, m + H_max <= last target slot of the session (389, or 209 on an early close), bar m present,5>= min_future_present of the H_max future slots present, s >= 1.6Branch A: last L_A (390) one-minute slots ending at slot m (crosses session boundaries, flagged).7Branch B: last L_B (390) completed 5-minute buckets (bucket j of session s is complete iff 5j+4 <= m).8Branch C: last L_C (260) completed 30-minute buckets (complete iff 30j+29 <= m).9Targets: y[h] = lc[m+h] - lc[prev present slot <= m+h-1] for h = 1..H_max with mask = present(m+h) (sum over present10h <= H equals the exact H-minute log return of the last observed bar). All inputs are scaled by the trailing sigma11(root mean square of present same-session 1-minute returns in the Branch-A window, floored), computed from the past only.12Grids are memory-mapped and opened lazily in each DataLoader worker (the index is pickle-light)."""13from __future__ import annotations1415import json16from dataclasses import dataclass17from pathlib import Path1819import numpy as np2021from .features import N_DAY, N_B5, N_B30, load_grid2223H_MAX = 6024SIGMA_FLOOR = 1e-525BAR_FEATURES = ["ret", "range", "upper", "lower", "body", "close_loc", "lv", "relvol", "missing", "boundary", "gap",26 "minute_frac", "time_to_close", "weekday", "early_close"]27F_BAR = len(BAR_FEATURES)28GLOBAL_FEATURES = ["log_sigma", "overnight_gap_scaled", "minute_frac", "time_to_close", "weekday", "early_close",29 "rv_30_rel", "rv_day_rel", "relvol_30"]30F_GLOBAL = len(GLOBAL_FEATURES)313233@dataclass34class WindowConfig:35 L_A: int = 39036 L_B: int = 39037 L_C: int = 26038 H: int = H_MAX39 m_min: int = 3040 min_future_present: float = 0.841 extended: bool = False424344class SymbolGrid:45 """Memory-mapped per-instrument grid with light derived metadata."""4647 def __init__(self, symbol: str, arrays: dict, meta: dict):48 self.symbol, self.arrays, self.meta = symbol, arrays, meta49 base = arrays["base"]; mask = arrays["mask"]50 self.S = base.shape[0]51 self.n_slots = np.asarray(meta["n_slots"])52 self.early = np.asarray(meta["early_close"], dtype=bool)53 wd = np.array([np.datetime64(d).astype("datetime64[D]").astype(int) for d in meta["session_dates"]])54 self.weekday = (wd + 3) % 7 # 1970-01-01 was a Thursday -> Monday = 055 self.gap = np.array([np.nan if g is None else g for g in meta["overnight_gap"]], dtype=np.float32)56 self.last_target_slot = np.where(self.early, 209, 389)57 # flat views (reshape of a memmap is a view, no copy)58 self.flat_base = base[:, :N_DAY, :].reshape(-1, 5)59 self.flat_mask = mask[:, :N_DAY].reshape(-1)60 self.flat_seas = arrays["seas"][:, :N_DAY].reshape(-1)61 self.flat_b5 = arrays["b5"].reshape(-1, 5); self.flat_m5 = arrays["m5"].reshape(-1)62 self.flat_b30 = arrays["b30"].reshape(-1, 5); self.flat_m30 = arrays["m30"].reshape(-1)63 self.flat_lc = arrays["flat_lc"]; self.prev_present = arrays["prev_present"]64 self.flat_seas_r2 = arrays["seas_r2"].reshape(-1)6566 @classmethod67 def load(cls, dir_: Path, mmap: bool = True) -> "SymbolGrid":68 arrays, meta = load_grid(Path(dir_), mmap=mmap)69 return cls(meta["symbol"], arrays, meta)707172def eligible_origins(g: SymbolGrid, cfg: WindowConfig, sessions: np.ndarray | None = None, stride: int = 1) -> np.ndarray:73 """(s, m) origins satisfying the eligibility rules (vectorised)."""74 mask = np.asarray(g.arrays["mask"][:, :N_DAY]) > 075 S = g.S76 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)77 if len(sess) == 0:78 return np.zeros((0, 2), dtype=np.int64)79 cs = np.concatenate([np.zeros((S, 1)), np.cumsum(mask, axis=1)], axis=1) # cs[s, k] = present in slots < k80 out = []81 for s in sess:82 last = int(g.last_target_slot[s])83 m = np.arange(cfg.m_min, last - cfg.H + 1, stride)84 if len(m) == 0:85 continue86 fut = (cs[s, m + 1 + cfg.H] - cs[s, m + 1]) / cfg.H87 ok = mask[s, m] & (fut >= cfg.min_future_present)88 if ok.any():89 out.append(np.column_stack([np.full(ok.sum(), s), m[ok]]))90 return np.concatenate(out).astype(np.int64) if out else np.zeros((0, 2), dtype=np.int64)919293def _bar_block(b: np.ndarray, m: np.ndarray, seas: np.ndarray | None, flat_idx: np.ndarray, per_day: int,94 sigma: float, scale: float, g: SymbolGrid) -> np.ndarray:95 L = len(m)96 f = np.zeros((L, F_BAR), dtype=np.float32)97 lc = np.where(m, b[:, 3], np.nan)98 # previous present close within the window (forward carry of the last observed close)99 idx = np.where(m, np.arange(L), -1)100 pp = np.maximum.accumulate(idx)101 prev = np.full(L, np.nan)102 prev[1:] = np.where(pp[:-1] >= 0, lc[np.clip(pp[:-1], 0, None)], np.nan)103 ret = np.where(m & ~np.isnan(prev), lc - prev, 0.0)104 s_eff = sigma * scale105 f[:, 0] = ret / s_eff106 with np.errstate(invalid="ignore"):107 rng = np.where(m, b[:, 1] - b[:, 2], 0.0)108 f[:, 1] = rng / s_eff109 f[:, 2] = np.where(m, b[:, 1] - np.maximum(b[:, 0], b[:, 3]), 0.0) / s_eff110 f[:, 3] = np.where(m, np.minimum(b[:, 0], b[:, 3]) - b[:, 2], 0.0) / s_eff111 f[:, 4] = np.where(m, b[:, 3] - b[:, 0], 0.0) / s_eff112 f[:, 5] = np.where(m & (rng > 0), (b[:, 3] - b[:, 2]) / np.where(rng > 0, rng, 1.0), 0.5) - 0.5113 f[:, 6] = np.where(m, b[:, 4], 0.0) / 10.0114 f[:, 7] = np.where(m & ~np.isnan(seas), b[:, 4] - seas, 0.0) if seas is not None else 0.0115 f[:, 8] = (~m).astype(np.float32)116 sess = flat_idx // per_day117 slot = flat_idx % per_day118 boundary = np.zeros(L, dtype=np.float32)119 boundary[1:] = (sess[1:] != sess[:-1]).astype(np.float32)120 f[:, 9] = boundary121 width = 390 // per_day122 gap_min = np.full(L, float(width))123 if L > 1:124 cross = sess[1:] != sess[:-1]125 # minutes from the previous row's bar open to this row's bar open across the session change126 rem_prev = (g.n_slots[sess[:-1]] - 1) - (slot[:-1] + 1) * width + 1127 gap_cross = width + rem_prev + (24 * 60 - 390) * np.maximum(1, sess[1:] - sess[:-1])128 gap_min[1:] = np.where(cross, gap_cross, width)129 f[:, 10] = np.log1p(gap_min) / 8.0130 f[:, 11] = (slot * width) / 389.0131 f[:, 12] = (g.n_slots[sess] - 1 - slot * width) / 389.0132 f[:, 13] = g.weekday[sess] / 4.0133 f[:, 14] = g.early[sess].astype(np.float32)134 return np.nan_to_num(f, nan=0.0, posinf=0.0, neginf=0.0)135136137def trailing_returns(g: SymbolGrid, s: int, m: int, L: int = 390) -> np.ndarray:138 """Same-session 1-minute returns between present bars inside the last L+1 slots ending at (s, m)."""139 end = s * N_DAY + m + 1140 lo = max(0, end - L - 1)141 idx = np.arange(lo, end)142 lc = np.asarray(g.flat_lc[idx], dtype=np.float64)143 ok = ~np.isnan(lc)144 if ok.sum() < 2:145 return np.zeros(0)146 r = np.diff(lc[ok]); sess = idx[ok] // N_DAY147 return r[sess[1:] == sess[:-1]]148149150def trailing_sigma(g: SymbolGrid, s: int, m: int, L: int = 390) -> tuple[float, np.ndarray]:151 r = trailing_returns(g, s, m, L)152 sigma = float(np.sqrt(np.mean(r ** 2))) if len(r) else SIGMA_FLOOR153 return max(sigma, SIGMA_FLOOR), r154155156def make_example(g: SymbolGrid, s: int, m: int, cfg: WindowConfig) -> dict[str, np.ndarray]:157 sigma, r_hist = trailing_sigma(g, s, m, cfg.L_A)158 endA = s * N_DAY + m + 1159 idxA = np.arange(endA - cfg.L_A, endA); validA = idxA >= 0; idxA_c = np.clip(idxA, 0, None)160 bA = np.asarray(g.flat_base[idxA_c]); mA = (np.asarray(g.flat_mask[idxA_c]) > 0) & validA161 xA = _bar_block(bA, mA, np.asarray(g.flat_seas[idxA_c]), idxA_c, N_DAY, sigma, 1.0, g)162 xA[~validA] = 0.0; xA[~validA, 8] = 1.0163 jB = (m + 1) // 5 - 1164 endB = s * N_B5 + jB + 1165 idxB = np.arange(endB - cfg.L_B, endB); validB = idxB >= 0; idxB_c = np.clip(idxB, 0, None)166 bB = np.asarray(g.flat_b5[idxB_c]); mB = (np.asarray(g.flat_m5[idxB_c]) > 0) & validB167 xB = _bar_block(bB, mB, None, idxB_c, N_B5, sigma, np.sqrt(5.0), g)168 xB[~validB] = 0.0; xB[~validB, 8] = 1.0169 jC = (m + 1) // 30 - 1170 endC = s * N_B30 + jC + 1171 idxC = np.arange(endC - cfg.L_C, endC); validC = idxC >= 0; idxC_c = np.clip(idxC, 0, None)172 bC = np.asarray(g.flat_b30[idxC_c]); mC = (np.asarray(g.flat_m30[idxC_c]) > 0) & validC173 xC = _bar_block(bC, mC, None, idxC_c, N_B30, sigma, np.sqrt(30.0), g)174 xC[~validC] = 0.0; xC[~validC, 8] = 1.0175 base_idx = s * N_DAY + m176 tgt_idx = base_idx + np.arange(1, cfg.H + 1)177 lc_t = np.asarray(g.flat_lc[tgt_idx], dtype=np.float64)178 prev_idx = np.asarray(g.prev_present[tgt_idx - 1])179 y_mask = ~np.isnan(lc_t) & (prev_idx >= 0)180 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)181 fut_slot = m + np.arange(1, cfg.H + 1)182 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),183 np.full(cfg.H, float(g.early[s]))], axis=1).astype(np.float32)184 rv30 = float(np.mean(r_hist[-30:] ** 2)) if len(r_hist) >= 5 else sigma ** 2185 rvday = float(np.mean(r_hist[-m:] ** 2)) if m >= 5 and len(r_hist) >= 5 else sigma ** 2186 relvol30 = float(np.mean(xA[-30:, 7]))187 gap = g.gap[s]188 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,189 g.weekday[s] / 4.0, float(g.early[s]), np.log(rv30 / sigma ** 2 + 1e-12), np.log(rvday / sigma ** 2 + 1e-12),190 relvol30], dtype=np.float32)191 return {"xA": xA, "mA": mA.astype(np.uint8), "xB": xB, "mB": mB.astype(np.uint8), "xC": xC, "mC": mC.astype(np.uint8),192 "y": y, "y_mask": y_mask.astype(np.uint8), "fut_cal": fut_cal, "glob": np.nan_to_num(glob),193 "sigma": np.float32(sigma), "origin": np.array([s, m], dtype=np.int64)}194195196class OriginIndex:197 """Flat index of (symbol, session, slot) origins across many grids. Holds grid directories; grids are opened198 lazily (memory-mapped) so the index can be pickled to DataLoader workers cheaply."""199200 def __init__(self, grid_dirs: list[Path], cfg: WindowConfig, session_sets: list[np.ndarray] | None = None,201 stride: int = 1, max_per_session: int | None = None, seed: int = 0, grids: list[SymbolGrid] | None = None):202 self.grid_dirs = [Path(d) for d in grid_dirs]203 self.cfg = cfg204 self._grids: dict[int, SymbolGrid] = {}205 if grids is not None:206 self._grids = dict(enumerate(grids))207 rng = np.random.default_rng(seed)208 rows = []209 for i in range(len(self.grid_dirs)):210 g = self.grid(i)211 o = eligible_origins(g, cfg, None if session_sets is None else session_sets[i], stride)212 if max_per_session is not None and len(o):213 # vectorised per-session subsampling: random key, sort by (session, key), keep the first k per session214 keys = rng.random(len(o))215 order = np.lexsort((keys, o[:, 0]))216 os_ = o[order]217 starts = np.r_[0, np.flatnonzero(np.diff(os_[:, 0])) + 1]218 sizes = np.diff(np.r_[starts, len(os_)])219 rank = np.arange(len(os_)) - np.repeat(starts, sizes)220 o = os_[rank < max_per_session]221 o = o[np.lexsort((o[:, 1], o[:, 0]))]222 if len(o):223 rows.append(np.column_stack([np.full(len(o), i), o]))224 self.rows = np.concatenate(rows).astype(np.int64) if rows else np.zeros((0, 3), dtype=np.int64)225 self.symbols = [g.symbol for g in (self.grid(i) for i in range(len(self.grid_dirs)))]226227 def grid(self, i: int) -> SymbolGrid:228 g = self._grids.get(i)229 if g is None:230 g = SymbolGrid.load(self.grid_dirs[i], mmap=True)231 self._grids[i] = g232 return g233234 @property235 def grids(self) -> list[SymbolGrid]:236 return [self.grid(i) for i in range(len(self.grid_dirs))]237238 def __getstate__(self):239 d = self.__dict__.copy(); d["_grids"] = {}240 return d241242 def __len__(self):243 return len(self.rows)244245 def example(self, k: int) -> dict[str, np.ndarray]:246 i, s, m = self.rows[k]247 ex = make_example(self.grid(int(i)), int(s), int(m), self.cfg)248 ex["sym"] = np.int64(i)249 return ex250251252def save_index(idx: OriginIndex, path: Path) -> None:253 np.save(path, idx.rows)254