1"""Torch dataset over an OriginIndex with feature-set masking (ablations) and a padded collate."""2from __future__ import annotations34import numpy as np5import torch6from torch.utils.data import Dataset78from ..data.windows import OriginIndex, BAR_FEATURES910FEATURE_SETS = {11 "full": None,12 "returns_only": ["ret", "missing", "boundary", "gap", "minute_frac", "time_to_close", "weekday", "early_close"],13 "ohlcv": [f for f in BAR_FEATURES if f != "relvol"],14 "time_only": ["missing", "boundary", "gap", "minute_frac", "time_to_close", "weekday", "early_close"],15}161718class OriginDataset(Dataset):19 def __init__(self, index: OriginIndex, feature_set: str = "full", shuffle_minutes: bool = False,20 shuffle_seed: int = 0):21 self.index, self.shuffle_minutes = index, shuffle_minutes22 keep = FEATURE_SETS[feature_set]23 self.fmask = None if keep is None else np.array([f in keep for f in BAR_FEATURES], dtype=np.float32)24 self.rng = np.random.default_rng(shuffle_seed)2526 def __len__(self):27 return len(self.index)2829 def __getitem__(self, k: int) -> dict[str, np.ndarray]:30 ex = self.index.example(k)31 if self.fmask is not None:32 for b in ("xA", "xB", "xC"):33 ex[b] = ex[b] * self.fmask34 if "ret" not in (FEATURE_SETS and []):35 pass36 if self.shuffle_minutes: # falsification: destroy temporal order inside the 1-minute window37 perm = self.rng.permutation(ex["xA"].shape[0])38 ex["xA"] = ex["xA"][perm]; ex["mA"] = ex["mA"][perm]39 return ex404142def collate(items: list[dict[str, np.ndarray]]) -> dict[str, torch.Tensor]:43 out = {}44 for k in items[0]:45 arr = np.stack([it[k] for it in items])46 out[k] = torch.from_numpy(arr)47 return out48