feat(model): extend TimeGAN with training loop, ckpt I/O, KL check, and generation API

Adds full wrapper (optimizers, ER pretrain, supervised, joint phases), checkpoint save/load, quick KL(spread) validation, and deterministic helpers. Integrates dataset batcher and utils (minmax, noise). Exposes encoder/recovery/generator/supervisor/discriminator and device/seed utilities.

Author: keys-i

recognition/TimeLOB_TimeGAN_49088276/src/helpers/args.py

@@ -68,6 +68,64 @@ def parse(self, argv: Optional[list | str]) -> Namespace:
 
         return ns
 
+class ModulesOptions:
+    """
+    Hyperparameters for modules & training. Designed to feel like an `opt` object.
+
+    Usage:
+        mods = ModulesOptions().parse(argv_after_flag)
+        # Access:
+        mods.batch_size, mods.seq_len, mods.z_dim, mods.hidden_dim, mods.num_layer,
+        mods.lr, mods.beta1, mods.w_gamma, mods.w_g
+    """
+    def __init__(self) -> None:
+        parser = ArgumentParser(
+            prog="timeganlob_modules",
+            description="Module/model hyperparameters and training weights.",
+        )
+        # Core shapes
+        parser.add_argument("--batch-size", type=int, default=128)
+        parser.add_argument("--seq-len", type=int, default=128,
+                            help="Sequence length (kept here for convenience to sync with data).")
+        parser.add_argument("--z-dim", type=int, default=40,
+                            help="Latent/input feature dim (e.g., LOB feature count).")
+        parser.add_argument("--hidden-dim", type=int, default=64,
+                            help="Module hidden size.")
+        parser.add_argument("--num-layer", type=int, default=3,
+                            help="Number of stacked layers per RNN/TCN block.")
+
+        # Optimizer
+        parser.add_argument("--lr", type=float, default=1e-4,
+                            help="Learning rate (generator/supervisor/discriminator if shared).")
+        parser.add_argument("--beta1", type=float, default=0.5,
+                            help="Adam beta1.")
+
+        # Loss weights
+        parser.add_argument("--w-gamma", type=float, default=1.0,
+                            help="Supervisor loss weight (γ).")
+        parser.add_argument("--w-g", type=float, default=1.0,
+                            help="Generator adversarial loss weight (g).")
+
+        self._parser = parser
+
+    def parse(self, argv: Optional[list | str]) -> Namespace:
+        m = self._parser.parse_args(argv)
+
+        # Provide both snake_case and "opt-like" names already as attributes
+        # (so downstream code can do opt.lr, opt.beta1, opt.w_gamma, opt.w_g).
+        ns = Namespace(
+            batch_size=m.batch_size,
+            seq_len=m.seq_len,
+            z_dim=m.z_dim,
+            hidden_dim=m.hidden_dim,
+            num_layer=m.num_layer,
+            lr=m.lr,
+            beta1=m.beta1,
+            w_gamma=m.w_gamma,
+            w_g=m.w_g,
+        )
+        return ns
+
 class Options:
     """
     Top-level options that *route* anything after `--dataset` to DatasetOptions.
@@ -92,6 +150,14 @@ def __init__(self) -> None:
                 "Example: --dataset --seq-len 256 --no-shuffle"
             ),
         )
+        parser.add_argument(
+            "--modules",
+            nargs=REMAINDER,
+            help=(
+                "All arguments following this flag are parsed by ModulesOptions. "
+                "Example: --modules --batch-size 256 --hidden-dim 128 --lr 3e-4"
+            ),
+        )
         self._parser = parser
 
     def parse(self, argv: Optional[list | str] = None) -> Namespace:

recognition/TimeLOB_TimeGAN_49088276/src/helpers/constants.py

@@ -3,7 +3,12 @@
 """
 from math import isclose
 from typing import Literal
+
 OUTPUT_DIR = "outs"
+WEIGHTS_DIR = "weights"
+DATA_DIR = "data"
+
+ORDERBOOK_FILENAME = "AMZN_2012-06-21_34200000_57600000_orderbook_10.csv"
 
 # Training hyperparameters for TimeGAN
 NUM_TRAINING_ITERATIONS = 25_000
@@ -16,8 +21,3 @@
 ), (
     f"TRAIN_TEST_SPLIT must sum to 1.0 (got {sum(TRAIN_TEST_SPLIT):.8f})"
 )
-
-DATA_DIR = "data"
-ORDERBOOK_FILENAME = "AMZN_2012-06-21_34200000_57600000_orderbook_10.csv"
-
-DATANAME = Literal["message", "orderbook"]