feat(dataset): add DataOptions CLI; robust split handling; logging; fix batch_generator

Introduce DataOptions wrapper with flags (--seq_len, --data_dir, --orderbook_filename, --no_shuffle, --keep_zero_rows, --splits, --log_level). Support ORDERBOOK_DEFAULT/SPLITS_DEFAULT fallbacks; accept proportions or cumulative cutoffs; replace prints with logging; add CLI entrypoint. Fix batch_generator index sampling and time=None handling; return constant T_mb; return windowed splits from load_data.

Author: keys-i

recognition/TimeLOB_TimeGAN_49088276/src/dataset.py

@@ -64,9 +64,9 @@ class DatasetConfig:
     """
     seq_len: int
     data_dir: Path = field(default_factory=lambda: Path(DATA_DIR))
-    filename: str = ORDERBOOK_FILENAME
+    orderbook_filename: str = ORDERBOOK_FILENAME
     splits: Tuple[float, float, float] = TRAIN_TEST_SPLIT
-    shuffle: bool = True
+    shuffle_windows: bool = True
     dtype: type = np.float32
     filter_zero_rows: bool = True
 
@@ -75,8 +75,8 @@ def from_namespace(cls, arg: Namespace) -> "DatasetConfig":
         return cls(
             seq_len=getattr(arg, "seq_len", 128),
             data_dir=Path(getattr(arg, "data_dir", DATA_DIR)),
-            filename=getattr(arg, "filename", ORDERBOOK_FILENAME),
-            shuffle=getattr(arg, "shuffle", True),
+            orderbook_filename=getattr(arg, "orderbook_filename", ORDERBOOK_FILENAME),
+            shuffle_windows=getattr(arg, "shuffle_windows", True),
             dtype=getattr(arg, "dtype", np.float32),
             filter_zero_rows=getattr(arg, "filter_zero_rows", True),
         )
@@ -119,7 +119,7 @@ def make_windows(
         Window the selected split into shape (num_windows, seq_len, num_features).
         """
         data = self._select_split(split)
-        return self._windowize(data, self.cfg.seq_len, self.cfg.shuffle)
+        return self._windowize(data, self.cfg.seq_len, self.cfg.shuffle_windows)
 
     def dataset_windowed(
             self
@@ -133,7 +133,7 @@ def dataset_windowed(
         return train_w, val_w, test_w
 
     def _read_raw(self) -> NDArray[np.int64]:
-        path = Path(self.cfg.data_dir, self.cfg.filename)
+        path = Path(self.cfg.data_dir, self.cfg.orderbook_filename)
         if not path.exists():
             msg = (
                 f"{path} not found.\n"
@@ -166,6 +166,7 @@ def _split_chronological(self) -> None:
         self._train = self._filtered[:t_cutoff]
         self._val = self._filtered[t_cutoff:v_cutoff]
         self._test = self._filtered[v_cutoff:]
+
         assert all(
             len(d) > 5 for d in (self._train, self._val, self._test)
         ), "Each split must have at least 5 windows."
@@ -186,7 +187,7 @@ def _windowize(
             self,
             data: NDArray[np.float32],
             seq_len: int,
-            shuffle: bool
+            shuffle_windows: bool
     ) -> NDArray[np.float32]:
         n_samples, n_features = data.shape
         n_windows = n_samples - seq_len + 1
@@ -196,7 +197,7 @@ def _windowize(
         out = np.empty((n_windows, seq_len, n_features), dtype=self.cfg.dtype)
         for i in range(n_windows):
             out[i] = data[i: i + seq_len]
-        if shuffle:
+        if shuffle_windows:
             np.random.shuffle(out)
         return out
 
@@ -217,13 +218,13 @@ def batch_generator(
     if `time` is None, uses a constant length equal to data.shape[1] (seq_len).
     """
     n = len(data)
-    idx = np.random.randint(n)[:batch_size]
+    idx = np.random.choice(n, size=batch_size, replace=True)
     data_mb = data[idx].astype(np.float32)
     if time is not None:
-        T_mb = np.full((batch_size,), data_mb.shape[1], dtype=np.int32)
+        t_mb = np.full((batch_size,), data_mb.shape[1], dtype=np.int32)
     else:
-        T_mb = time[idx].astype(np.int32)
-    return data_mb, T_mb
+        t_mb = time[idx].astype(np.int32)
+    return data_mb, t_mb
 
 
 def load_data(arg: Namespace) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:

recognition/TimeLOB_TimeGAN_49088276/src/helpers/arg2.py

@@ -0,0 +1,69 @@
+"""
+Options for the entire model
+"""
+from __future__ import annotations
+
+from argparse import ArgumentParser, Namespace
+from typing import Optional
+
+import numpy as np
+
+from src.helpers.constants import DATA_DIR, TRAIN_TEST_SPLIT, ORDERBOOK_FILENAME
+
+try:
+    # tolerate alternates if present in your helpers
+    from src.helpers.constants import ORDERBOOK_FILENAME as _OB_ALT
+    ORDERBOOK_DEFAULT = _OB_ALT
+except Exception:
+    ORDERBOOK_DEFAULT = ORDERBOOK_FILENAME
+
+class DataOptions:
+    """
+    Thin wrapper around argparse that produces a Namespace suitable for DatasetConfig.
+    Usage:
+        opts = DataOptions().parse()
+        train_w, val_w, test_w = load_data(opts)
+    """
+
+    def __init__(self) -> None:
+        parser = ArgumentParser(
+            prog="timeganlob_dataset",
+            description="Lightweight LOBSTER preprocessing + MinMax scaling",
+        )
+        parser.add_argument("--seq-len", type=int, default=128)
+        parser.add_argument("--data_dir", type=str, default=str(DATA_DIR))
+        parser.add_argument("--orderbook_filename", type=str, default=ORDERBOOK_FILENAME)
+        parser.add_argument(
+            "--no-shuffle",
+            action="store_true",
+            help="Disable shuffling of windowed sequences"
+        )
+        parser.add_argument(
+            "--keep_zero_rows",
+            action="store_true",
+            help="Do NOT filter rows containing zeros."
+        )
+        parser.add_argument(
+            "--splits",
+            type=float,
+            nargs=3,
+            metavar=("TRAIN", "VAL", "TEST"),
+            help="Either proportions that sum to ~1.0 or cumulative cutoffs (e.g., 0.6 0.8 1.0).",
+            default=None,
+        )
+        self._parser = parser
+
+    def parse(self, argv: Optional[list | str]) -> Namespace:
+        args = self._parser.parse_args(argv)
+
+        ns = Namespace(
+            seq_len=args.seq_len,
+            data_dir=args.data_dir,
+            orderbook_filename=args.orderbook_filename,
+            splits=tuple(args.splits) if args.splits is not None else TRAIN_TEST_SPLIT,
+            shuffle_windows=not args.no_shuffle,
+            dtype=np.float32,
+            keep_zero_rows=not args.keep_zero_rows,
+        )
+
+        return ns