Add inference code, config, and technical report

config/folds.json

@@ -0,0 +1,1057 @@
+{
+  "dataset_sha256": "79c14580e40baa6c0d08faf84d3b754e92be28aaf80fc6be4a3b341ea35a0544",
+  "final_holdout_fold": 4,
+  "final_holdout_seq_ids": [
+    1,
+    13,
+    20,
+    21,
+    34,
+    43,
+    48,
+    49,
+    50,
+    52,
+    53,
+    54,
+    58,
+    71,
+    80,
+    87,
+    88,
+    91,
+    99,
+    102,
+    105,
+    106,
+    121,
+    130,
+    134,
+    149,
+    151,
+    160,
+    161,
+    166,
+    169,
+    174,
+    187,
+    188,
+    189,
+    190,
+    191,
+    201,
+    205,
+    214,
+    217,
+    235,
+    241,
+    243,
+    251,
+    252,
+    257,
+    259,
+    263,
+    264,
+    269,
+    270,
+    273,
+    276,
+    293,
+    303,
+    306,
+    308,
+    309,
+    313,
+    315,
+    319,
+    328,
+    330,
+    339,
+    343,
+    344,
+    345,
+    350,
+    359,
+    363,
+    366,
+    372,
+    385,
+    387,
+    389,
+    401,
+    413,
+    419,
+    427,
+    430,
+    435,
+    443,
+    445,
+    455,
+    458,
+    459,
+    463,
+    466,
+    470,
+    471,
+    472,
+    474,
+    475,
+    488,
+    490,
+    491,
+    498,
+    500,
+    506,
+    510,
+    514,
+    516
+  ],
+  "folds": {
+    "0": [
+      0,
+      2,
+      9,
+      10,
+      11,
+      18,
+      22,
+      30,
+      33,
+      39,
+      46,
+      55,
+      63,
+      68,
+      69,
+      70,
+      72,
+      73,
+      76,
+      77,
+      78,
+      79,
+      82,
+      84,
+      90,
+      93,
+      101,
+      104,
+      124,
+      131,
+      132,
+      137,
+      140,
+      148,
+      153,
+      155,
+      167,
+      172,
+      176,
+      177,
+      184,
+      194,
+      195,
+      203,
+      204,
+      208,
+      209,
+      220,
+      222,
+      225,
+      227,
+      231,
+      237,
+      238,
+      244,
+      247,
+      249,
+      250,
+      253,
+      275,
+      277,
+      281,
+      287,
+      289,
+      296,
+      298,
+      304,
+      317,
+      318,
+      320,
+      323,
+      324,
+      325,
+      331,
+      341,
+      342,
+      353,
+      357,
+      362,
+      382,
+      398,
+      400,
+      403,
+      405,
+      426,
+      429,
+      431,
+      441,
+      447,
+      451,
+      454,
+      465,
+      469,
+      477,
+      487,
+      489,
+      501,
+      503,
+      504,
+      507,
+      508,
+      509,
+      512,
+      513
+    ],
+    "1": [
+      3,
+      5,
+      7,
+      15,
+      16,
+      17,
+      19,
+      23,
+      24,
+      25,
+      26,
+      29,
+      31,
+      37,
+      42,
+      45,
+      56,
+      57,
+      60,
+      66,
+      67,
+      75,
+      86,
+      94,
+      108,
+      110,
+      113,
+      114,
+      116,
+      117,
+      118,
+      126,
+      141,
+      144,
+      154,
+      157,
+      168,
+      173,
+      175,
+      180,
+      181,
+      182,
+      185,
+      192,
+      210,
+      211,
+      218,
+      228,
+      234,
+      255,
+      265,
+      266,
+      268,
+      271,
+      272,
+      274,
+      278,
+      284,
+      286,
+      299,
+      302,
+      311,
+      312,
+      321,
+      329,
+      332,
+      334,
+      340,
+      349,
+      352,
+      356,
+      367,
+      368,
+      373,
+      376,
+      380,
+      384,
+      391,
+      393,
+      396,
+      397,
+      408,
+      409,
+      412,
+      416,
+      420,
+      421,
+      422,
+      433,
+      438,
+      446,
+      452,
+      456,
+      457,
+      467,
+      468,
+      479,
+      483,
+      485,
+      496,
+      497,
+      505,
+      511,
+      515
+    ],
+    "2": [
+      6,
+      36,
+      38,
+      59,
+      74,
+      81,
+      83,
+      89,
+      92,
+      96,
+      97,
+      103,
+      109,
+      111,
+      112,
+      119,
+      122,
+      123,
+      125,
+      129,
+      139,
+      143,
+      145,
+      146,
+      147,
+      150,
+      152,
+      158,
+      163,
+      165,
+      179,
+      183,
+      193,
+      196,
+      198,
+      199,
+      202,
+      219,
+      223,
+      229,
+      233,
+      239,
+      245,
+      246,
+      248,
+      261,
+      262,
+      280,
+      290,
+      291,
+      294,
+      297,
+      301,
+      305,
+      307,
+      310,
+      314,
+      316,
+      322,
+      335,
+      336,
+      338,
+      346,
+      347,
+      348,
+      354,
+      360,
+      361,
+      364,
+      365,
+      370,
+      371,
+      375,
+      377,
+      386,
+      388,
+      392,
+      394,
+      402,
+      404,
+      406,
+      407,
+      411,
+      417,
+      418,
+      428,
+      432,
+      440,
+      442,
+      450,
+      453,
+      461,
+      464,
+      473,
+      481,
+      482,
+      484,
+      486,
+      492,
+      493,
+      495,
+      499,
+      502
+    ],
+    "3": [
+      4,
+      8,
+      12,
+      14,
+      27,
+      28,
+      32,
+      35,
+      40,
+      41,
+      44,
+      47,
+      51,
+      61,
+      62,
+      64,
+      65,
+      85,
+      95,
+      98,
+      100,
+      107,
+      115,
+      120,
+      127,
+      128,
+      133,
+      135,
+      136,
+      138,
+      142,
+      156,
+      159,
+      162,
+      164,
+      170,
+      171,
+      178,
+      186,
+      197,
+      200,
+      206,
+      207,
+      212,
+      213,
+      215,
+      216,
+      221,
+      224,
+      226,
+      230,
+      232,
+      236,
+      240,
+      242,
+      254,
+      256,
+      258,
+      260,
+      267,
+      279,
+      282,
+      283,
+      285,
+      288,
+      292,
+      295,
+      300,
+      326,
+      327,
+      333,
+      337,
+      351,
+      355,
+      358,
+      369,
+      374,
+      378,
+      379,
+      381,
+      383,
+      390,
+      395,
+      399,
+      410,
+      414,
+      415,
+      423,
+      424,
+      425,
+      434,
+      436,
+      437,
+      439,
+      444,
+      448,
+      449,
+      460,
+      462,
+      476,
+      478,
+      480,
+      494
+    ],
+    "4": [
+      1,
+      13,
+      20,
+      21,
+      34,
+      43,
+      48,
+      49,
+      50,
+      52,
+      53,
+      54,
+      58,
+      71,
+      80,
+      87,
+      88,
+      91,
+      99,
+      102,
+      105,
+      106,
+      121,
+      130,
+      134,
+      149,
+      151,
+      160,
+      161,
+      166,
+      169,
+      174,
+      187,
+      188,
+      189,
+      190,
+      191,
+      201,
+      205,
+      214,
+      217,
+      235,
+      241,
+      243,
+      251,
+      252,
+      257,
+      259,
+      263,
+      264,
+      269,
+      270,
+      273,
+      276,
+      293,
+      303,
+      306,
+      308,
+      309,
+      313,
+      315,
+      319,
+      328,
+      330,
+      339,
+      343,
+      344,
+      345,
+      350,
+      359,
+      363,
+      366,
+      372,
+      385,
+      387,
+      389,
+      401,
+      413,
+      419,
+      427,
+      430,
+      435,
+      443,
+      445,
+      455,
+      458,
+      459,
+      463,
+      466,
+      470,
+      471,
+      472,
+      474,
+      475,
+      488,
+      490,
+      491,
+      498,
+      500,
+      506,
+      510,
+      514,
+      516
+    ]
+  },
+  "git_commit": "cf46f28c2a9df775367b428eedd3efd5a2f68399",
+  "n_folds": 5,
+  "seed": 42,
+  "train_dev_seq_ids": [
+    0,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    22,
+    23,
+    24,
+    25,
+    26,
+    27,
+    28,
+    29,
+    30,
+    31,
+    32,
+    33,
+    35,
+    36,
+    37,
+    38,
+    39,
+    40,
+    41,
+    42,
+    44,
+    45,
+    46,
+    47,
+    51,
+    55,
+    56,
+    57,
+    59,
+    60,
+    61,
+    62,
+    63,
+    64,
+    65,
+    66,
+    67,
+    68,
+    69,
+    70,
+    72,
+    73,
+    74,
+    75,
+    76,
+    77,
+    78,
+    79,
+    81,
+    82,
+    83,
+    84,
+    85,
+    86,
+    89,
+    90,
+    92,
+    93,
+    94,
+    95,
+    96,
+    97,
+    98,
+    100,
+    101,
+    103,
+    104,
+    107,
+    108,
+    109,
+    110,
+    111,
+    112,
+    113,
+    114,
+    115,
+    116,
+    117,
+    118,
+    119,
+    120,
+    122,
+    123,
+    124,
+    125,
+    126,
+    127,
+    128,
+    129,
+    131,
+    132,
+    133,
+    135,
+    136,
+    137,
+    138,
+    139,
+    140,
+    141,
+    142,
+    143,
+    144,
+    145,
+    146,
+    147,
+    148,
+    150,
+    152,
+    153,
+    154,
+    155,
+    156,
+    157,
+    158,
+    159,
+    162,
+    163,
+    164,
+    165,
+    167,
+    168,
+    170,
+    171,
+    172,
+    173,
+    175,
+    176,
+    177,
+    178,
+    179,
+    180,
+    181,
+    182,
+    183,
+    184,
+    185,
+    186,
+    192,
+    193,
+    194,
+    195,
+    196,
+    197,
+    198,
+    199,
+    200,
+    202,
+    203,
+    204,
+    206,
+    207,
+    208,
+    209,
+    210,
+    211,
+    212,
+    213,
+    215,
+    216,
+    218,
+    219,
+    220,
+    221,
+    222,
+    223,
+    224,
+    225,
+    226,
+    227,
+    228,
+    229,
+    230,
+    231,
+    232,
+    233,
+    234,
+    236,
+    237,
+    238,
+    239,
+    240,
+    242,
+    244,
+    245,
+    246,
+    247,
+    248,
+    249,
+    250,
+    253,
+    254,
+    255,
+    256,
+    258,
+    260,
+    261,
+    262,
+    265,
+    266,
+    267,
+    268,
+    271,
+    272,
+    274,
+    275,
+    277,
+    278,
+    279,
+    280,
+    281,
+    282,
+    283,
+    284,
+    285,
+    286,
+    287,
+    288,
+    289,
+    290,
+    291,
+    292,
+    294,
+    295,
+    296,
+    297,
+    298,
+    299,
+    300,
+    301,
+    302,
+    304,
+    305,
+    307,
+    310,
+    311,
+    312,
+    314,
+    316,
+    317,
+    318,
+    320,
+    321,
+    322,
+    323,
+    324,
+    325,
+    326,
+    327,
+    329,
+    331,
+    332,
+    333,
+    334,
+    335,
+    336,
+    337,
+    338,
+    340,
+    341,
+    342,
+    346,
+    347,
+    348,
+    349,
+    351,
+    352,
+    353,
+    354,
+    355,
+    356,
+    357,
+    358,
+    360,
+    361,
+    362,
+    364,
+    365,
+    367,
+    368,
+    369,
+    370,
+    371,
+    373,
+    374,
+    375,
+    376,
+    377,
+    378,
+    379,
+    380,
+    381,
+    382,
+    383,
+    384,
+    386,
+    388,
+    390,
+    391,
+    392,
+    393,
+    394,
+    395,
+    396,
+    397,
+    398,
+    399,
+    400,
+    402,
+    403,
+    404,
+    405,
+    406,
+    407,
+    408,
+    409,
+    410,
+    411,
+    412,
+    414,
+    415,
+    416,
+    417,
+    418,
+    420,
+    421,
+    422,
+    423,
+    424,
+    425,
+    426,
+    428,
+    429,
+    431,
+    432,
+    433,
+    434,
+    436,
+    437,
+    438,
+    439,
+    440,
+    441,
+    442,
+    444,
+    446,
+    447,
+    448,
+    449,
+    450,
+    451,
+    452,
+    453,
+    454,
+    456,
+    457,
+    460,
+    461,
+    462,
+    464,
+    465,
+    467,
+    468,
+    469,
+    473,
+    476,
+    477,
+    478,
+    479,
+    480,
+    481,
+    482,
+    483,
+    484,
+    485,
+    486,
+    487,
+    489,
+    492,
+    493,
+    494,
+    495,
+    496,
+    497,
+    499,
+    501,
+    502,
+    503,
+    504,
+    505,
+    507,
+    508,
+    509,
+    511,
+    512,
+    513,
+    515
+  ]
+}

reports/phase2_technical_report.md

@@ -0,0 +1,195 @@
+# Wunder Fund RNN Challenge — Technical Report
+
+**Task.** Predict the next 32-dimensional market state from a sequence of prior
+states, scored online (`PredictionModel.predict` one row at a time) by mean
+per-feature R². Inference budget: 1 CPU core, 16 GB RAM, no GPU, 60 minutes,
+deterministic. The competition is finished; this is a portfolio reconstruction
+whose target is the winner's `0.3964` R².
+
+---
+
+## 1. Executive summary
+
+The solution is a **causal unidirectional GRU** that predicts the next state at
+every step, carrying its hidden state across the sequence so each online
+prediction is O(1). It is trained on full 1000-step sequences with
+back-prop-through-time and a loss masked to the scored steps (current steps
+100–998 → targets 101–999). A multi-seed / multi-architecture ensemble is
+averaged per feature.
+
+Headline result (leaderboard-comparable, 5-fold grouped CV by `seq_ix`):
+
+| Model | Dev CV mean R² | Notes |
+| --- | ---: | --- |
+| Persistence (predict current) | −0.37 | sanity baseline |
+| EWMA α=0.10 | 0.217 | best stateful linear baseline |
+| Ridge causal (640 feats) | 0.327 | at the linear ceiling |
+| **Single causal GRU (d256)** | **0.389** | per-fold 0.391 / 0.394 / 0.407 / 0.363 |
+| **GRU ensemble (multi-seed/arch)** | _finalized below_ | target ≥ 0.396 |
+
+The single GRU already lands in the competition's top-10 band (0.386–0.392).
+Inference costs ≈ **0.7 ms/prediction** single-threaded, ~1 min for a full
+holdout, far inside the 60-minute budget.
+
+---
+
+## 2. Data analysis (what the data dictated)
+
+Measured with `scripts/analyze_data.py` on `data/raw/train.parquet`
+(517 sequences × 1000 steps × 32 features, columns named `"0"`…`"31"`):
+
+- **Pre-whitened features.** Global mean ≈ 0, std ≈ 1, per-feature std ∈
+  [0.985, 1.011], excess kurtosis ≈ 0, hard-clipped at ±5.199. A
+  Gaussianizing/quantile transform was applied → no heavy tails, MSE is a
+  well-behaved objective, and no bespoke normalization is needed.
+- **Mean-reverting, not a random walk.** Persistence R² ≈ −0.37 and
+  `var(next−current)/var(next) = 1.39` (differencing makes it *worse*), so the
+  model predicts the **level**, not a delta. Lag-1 autocorrelation ≈ 0.14.
+- **The linear ceiling is ≈ 0.32.** VAR(p) Ridge saturates: p=1→0.23, p=2→0.28,
+  p=4→0.30, p=8→0.31, p=16→0.32. The Ridge causal baseline (0.327) is already
+  at this ceiling.
+- **Nonlinearity is the lever (~+0.03).** Histogram gradient boosting on VAR(8)
+  features beats Ridge on every feature by +0.015…+0.055, with the largest
+  gains on the linearly-hardest features.
+
+Conclusion: linear models top out near 0.32; the ~0.07 gap to the winner is
+**nonlinear temporal structure**, which a sequence model is built to capture.
+This is why the primary model is a GRU rather than more feature engineering.
+
+---
+
+## 3. Architecture
+
+`src/models/sequence_models.py::CausalGRUForecaster`
+
+```
+state_t (32) ──► Linear(32→d) ─► LayerNorm ─► GRU(d, n_layers, unidirectional)
+                                                      │ h_t
+                                                      ▼
+                          LayerNorm ─► Linear(d→d) ─► GELU ─► Dropout ─► Linear(d→32) ─► pred of state_{t+1}
+```
+
+- **Unidirectional & causal** — uses only past/current states, matching the
+  online task. A bidirectional model would leak the future and is invalid here.
+- **Stateful O(1) inference** — `forward(x, h0)` and a one-step `step(state, h)`
+  share the same recurrence, so the batched training forward and the row-by-row
+  online replay are numerically identical (verified to the digit and in
+  `tests/test_sequence_model.py`).
+- Default size d_model=256, 2 layers (~0.87 M params); diversity members use
+  d_model=384 and a 3-layer variant.
+
+---
+
+## 4. Training strategy
+
+`src/training/sequence_trainer.py`
+
+- **Full-sequence BPTT.** Each 1000-step sequence is processed end-to-end; the
+  loss is a **masked MSE** over the scored steps only (`need_prediction`,
+  i.e. current steps 100–998 against the next state). Because targets are
+  unit-variance, MSE ≈ (1 − R²), so the objective directly tracks the metric.
+- **Optimizer/schedule.** AdamW (lr 1.5e-3, weight decay 1e-4), OneCycle cosine
+  schedule with 10% warmup, gradient clipping at 1.0.
+- **Early stopping** on batched validation R² (identical to the stepwise score),
+  patience 14, typically converging by epoch 55–65.
+- **Reproducibility.** Global seeds for Python/NumPy/PyTorch; deterministic
+  algorithms enabled; folds are fixed in `config/folds.json`.
+- **Hardware.** Trained on an NVIDIA RTX 3050 Laptop (4 GB) at ~1.8 s/epoch
+  (~2 min/run); a CPU-only path (Ryzen 7 5800H, ~24 s/epoch) reproduces the
+  same scores within ±0.0003. Submission inference is CPU-only by design.
+
+---
+
+## 5. Validation protocol
+
+`src/evaluation/stepwise.py` + `config/folds.json`
+
+- **Sequence-grouped** 5-fold split by `seq_ix` (seed 42); fold 4 is an
+  untouched final holdout, folds 0–3 are development. No row from a sequence
+  ever appears in both train and validation.
+- The **stepwise scorer** replays `PredictionModel.predict` row by row and never
+  scores across `seq_ix` boundaries, exactly mirroring the official scorer.
+- The **headline metric is the dev CV mean** across folds 0–3, not a single
+  fold: individual 103-sequence folds are high-variance (single-GRU range
+  0.363–0.407), whereas the hidden test (~517 sequences) averages over regimes.
+  Fold 4 is reported separately as a conservative point estimate (it is the
+  hardest fold; every dev fold scores higher).
+
+---
+
+## 6. Results
+
+Cross-validated R² (sequence-grouped folds 0–3), simple-average ensembles:
+
+| Ensemble | Members | CV mean R² |
+| --- | ---: | ---: |
+| Single GRU d256 (avg of 3 seeds) | 1 | 0.389 |
+| 3× GRU d256 | 3 | 0.3915 |
+| **3× GRU d256 + 2× GRU d384 (packaged)** | **5** | **0.3922** |
+| + 2× LSTM | 7 | 0.3925 |
+
+- Per-feature non-negative stacking matched simple averaging (0.3921 vs 0.3922)
+  — expected for a same-family ensemble, so the packaged model averages equally.
+- **Untouched fold-4 holdout** of the 5-member ensemble (fold models that never
+  saw fold 4): **0.3498** — the conservative single-fold figure; fold 4 is the
+  hardest fold (all dev folds score higher), which is why the CV mean is the
+  headline.
+- **Single-core inference:** 4.17 ms/prediction for the 5-member ensemble →
+  ≈ **32 minutes** for the full ~517-sequence test on one CPU core (limit: 60).
+- **Deterministic:** identical predictions across repeated runs (max abs diff 0).
+
+For reference, the competition's public leaderboard #1 was 0.3920 and the finals
+winner 0.3964; the cross-validated 0.3922 matches the public top.
+
+### Key boosters (most score per unit effort)
+1. **Causal GRU over linear/tree models** — moved 0.327 → 0.389 CV (+0.062), the
+   single biggest jump; it captures the nonlinear temporal structure linear
+   models cannot.
+2. **Predicting the level on whitened inputs** — differencing or persistence is
+   actively harmful here (negative R²); letting the GRU learn the mean-reverting
+   map directly is the correct framing.
+3. **Full-sequence stateful modeling** — carrying hidden state from step 0
+   matches train and inference exactly and makes online inference O(1).
+4. **Multi-seed + width-diversity ensembling** — +0.003 over a single GRU
+   (0.389 → 0.3922); the largest ensemble gains came from combining d256 and
+   d384 seeds rather than adding more of the same.
+
+## 7. What didn't work / didn't help
+- **More training data had little effect** (310 → 352 sequences: +0.003 on the
+  holdout). The dev→holdout gap is fold *difficulty*, not data quantity.
+- **Bagging same-config fold models** barely moved the holdout — correlated
+  members; diversity (seeds/architectures) is what helps.
+- **Heavy feature engineering / GBDT as the primary model** — capped near the
+  linear-plus-small-nonlinear ceiling (~0.34); kept only as an optional
+  diversity member, not the backbone.
+- **Differencing / momentum targets** — worse than predicting the level.
+- **Causal TCN (dilated convolutions)** — trained as a decorrelated member but
+  scored lower (≈0.37–0.39, fold-dependent) and was far slower at inference (a
+  ~250-step receptive field re-run each step). Doubly disqualifying, dropped.
+  (A real bug was caught here first: a time-mixing GroupNorm leaked the future;
+  fixed with a per-position LayerNorm — see `tests`/the equivalence check.)
+- **GBDT (HistGradientBoosting) as a stacking member** — strongly decorrelated
+  (≈0.37) but added only ~+0.0002 to the ensemble while being slow to train OOF;
+  not worth the inference complexity. Kept the ensemble pure-GRU.
+- **Adding more ensemble members past 5** — LSTM seeds lifted CV by only +0.0003
+  while increasing single-core inference cost; capped at 5 for budget headroom.
+- The earlier repo's "CMDMamba ensemble" was aspirational (no runnable
+  artifacts) and based on an unverifiable reference; it was discarded in favor
+  of a model justified by the measured data structure.
+
+## 8. Reproduction
+
+```bash
+# environment (CPU submission-faithful)
+py -3.11 -m venv .venv && .venv/Scripts/python -m pip install -r requirements.txt
+# optional GPU training env
+py -3.11 -m venv .venv-gpu && .venv-gpu/Scripts/python -m pip install -r requirements-train-gpu.txt
+
+python scripts/analyze_data.py                      # data diagnostics
+python scripts/create_folds.py --output config/folds.json
+python scripts/train_sequence.py --run-name gru --dev-fold 0 --device cuda   # one fold
+python scripts/cv_summary.py --runs <run...>        # ensemble CV mean
+python scripts/evaluate_ensemble.py --split holdout --gru <ckpt...>          # holdout
+python scripts/create_submission.py --output submissions/submission.zip
+python -m pytest tests -q
+```

requirements.txt

@@ -0,0 +1,36 @@
+# Wunder Fund RNN Challenge - Dependencies
+
+# Core
+numpy==1.26.4
+pandas==2.2.3
+pyarrow==18.1.0
+scikit-learn==1.5.2
+scipy==1.14.1
+joblib==1.4.2
+tqdm==4.66.5
+
+# Deep learning (CPU)
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.9.1
+einops==0.8.0
+
+# Tree-based models
+lightgbm==4.3.0
+xgboost==2.1.2
+
+# ONNX optimization
+onnx==1.17.0
+onnxruntime==1.20.1
+
+# Configuration
+pyyaml==6.0.2
+
+# Experiment tracking (optional)
+wandb==0.18.7
+
+# Hyperparameter optimization (optional)
+optuna==4.1.0
+
+# Testing
+pytest==8.3.4
+pytest-cov==6.0.0

solution.py

@@ -0,0 +1,198 @@
+"""solution.py — Competition submission for the Wunder Fund RNN Challenge.
+
+Root entry point for the scorer. Defines PredictionModel with a predict()
+method that receives DataPoint objects one at a time and predicts the *next*
+state when need_prediction is True.
+
+Primary backend: an ensemble of causal GRU forecasters (carrying hidden state
+across steps for O(1)/prediction online inference), optionally blended per
+feature with tabular models. Falls back to joblib tabular models, then a
+momentum baseline (only when ALLOW_FALLBACK=1).
+"""
+
+import os
+
+# Windows: torch and numpy/sklearn each ship an OpenMP runtime; importing torch
+# before the MKL-backed stack avoids a c10.dll init failure. Also keep the
+# submission single-core friendly by default (the scorer grants 1 CPU).
+os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
+_THREADS = os.environ.get("WUNDER_THREADS", "1")
+os.environ.setdefault("OMP_NUM_THREADS", _THREADS)
+
+try:
+    import torch
+    torch.set_num_threads(max(1, int(_THREADS)))
+    _HAS_TORCH = True
+except Exception:
+    _HAS_TORCH = False
+
+import json
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+_ROOT = Path(__file__).resolve().parent
+import sys
+if str(_ROOT) not in sys.path:
+    sys.path.insert(0, str(_ROOT))
+
+
+class DataPoint:
+    """Mirror of competition_package.utils.DataPoint for standalone use."""
+
+    def __init__(self, seq_ix, step_in_seq, need_prediction, state):
+        self.seq_ix = seq_ix
+        self.step_in_seq = step_in_seq
+        self.need_prediction = need_prediction
+        self.state = state
+
+
+class PredictionModel:
+    """Competition PredictionModel. Loads the best available backend."""
+
+    def __init__(self):
+        self.n_features = 32
+        self.backend = "fallback"
+        self.predictor = None  # object exposing predict(data_point)
+        self._current_seq = None
+        self._prev_state = None
+        self._cur_state = None
+
+        self._try_load_gru()
+        if self.predictor is None:
+            self._try_load_onnx()
+        if self.predictor is None:
+            self._try_load_tabular()
+        if self.predictor is None:
+            self._try_load_pytorch()
+        if self.predictor is None and os.environ.get("ALLOW_FALLBACK") != "1":
+            raise RuntimeError(
+                "No trained model artifacts found. Set ALLOW_FALLBACK=1 only for "
+                "local smoke tests; real submissions must package models."
+            )
+
+    # ------------------------------------------------------------------
+    # Backends
+    # ------------------------------------------------------------------
+    def _try_load_gru(self):
+        """Primary: GRU ensemble (+ optional tabular members) under models/submission."""
+        if not _HAS_TORCH:
+            return
+        try:
+            from src.models.ensemble_predictor import EnsemblePredictionModel
+            from src.models.sequence_inference import (
+                GRUStatefulPredictionModel,
+                load_gru_checkpoint,
+            )
+
+            sub = _ROOT / "models" / "submission"
+            gru_paths = sorted((sub / "gru").glob("*.pt")) if (sub / "gru").exists() else []
+            if not gru_paths:
+                return
+
+            members, names = [], []
+            for p in gru_paths:
+                model = load_gru_checkpoint(str(p))
+                members.append(GRUStatefulPredictionModel(model))
+                names.append(p.stem)
+
+            # optional tabular members
+            tab_dir = sub / "tabular"
+            if tab_dir.exists():
+                import joblib
+                from src.models.baselines import TabularStatefulPredictionModel
+                for jb in sorted(tab_dir.glob("*/model.joblib")):
+                    manifest = jb.parent / "model_manifest.json"
+                    schema = "compact_v1"
+                    if manifest.exists():
+                        try:
+                            schema = json.loads(manifest.read_text()).get("feature_schema", schema)
+                        except Exception:
+                            pass
+                    members.append(TabularStatefulPredictionModel(joblib.load(str(jb)), feature_schema=schema))
+                    names.append(jb.parent.name)
+
+            weights = None
+            wpath = sub / "blend_weights.json"
+            if wpath.exists():
+                try:
+                    payload = json.loads(wpath.read_text())
+                    w = np.asarray(payload["per_feature_weights"], dtype=np.float32)
+                    if w.shape == (self.n_features, len(members)):
+                        weights = w
+                except Exception:
+                    weights = None
+
+            self.predictor = EnsemblePredictionModel(members, weights=weights)
+            self.backend = "gru_ensemble"
+        except Exception:
+            self.predictor = None
+
+    def _try_load_onnx(self):
+        """ONNX ensemble (kept for completeness; not the primary path)."""
+        return
+
+    def _try_load_tabular(self):
+        """Joblib tabular models discovered under models/ (legacy/baseline path)."""
+        try:
+            import joblib
+            from src.models.baselines import TabularStatefulPredictionModel
+            from src.models.ensemble_predictor import EnsemblePredictionModel
+
+            model_dir = _ROOT / "models"
+            joblib_files = sorted(model_dir.glob("phase2_tabular/**/model.joblib"))
+            if not joblib_files:
+                return
+            members = []
+            for path in joblib_files:
+                manifest = path.parent / "model_manifest.json"
+                schema = "compact_v1"
+                if manifest.exists():
+                    try:
+                        schema = json.loads(manifest.read_text()).get("feature_schema", schema)
+                    except Exception:
+                        pass
+                members.append(TabularStatefulPredictionModel(joblib.load(str(path)), feature_schema=schema))
+            if members:
+                self.predictor = EnsemblePredictionModel(members)
+                self.backend = "tabular"
+        except Exception:
+            self.predictor = None
+
+    def _try_load_pytorch(self):
+        """Legacy PyTorch checkpoint path (kept for test compatibility)."""
+        return
+
+    # ------------------------------------------------------------------
+    # Predict
+    # ------------------------------------------------------------------
+    def predict(self, data_point) -> Optional[np.ndarray]:
+        if self.predictor is not None:
+            return self.predictor.predict(data_point)
+        return self._predict_fallback(data_point)
+
+    def _predict_fallback(self, data_point) -> Optional[np.ndarray]:
+        """Momentum baseline; only reachable with ALLOW_FALLBACK=1."""
+        if self._current_seq != data_point.seq_ix:
+            self._current_seq = data_point.seq_ix
+            self._prev_state = None
+            self._cur_state = None
+        self._prev_state = self._cur_state
+        self._cur_state = np.asarray(data_point.state, dtype=np.float32)
+        if not data_point.need_prediction:
+            return None
+        if self._prev_state is None:
+            return self._cur_state.astype(np.float32, copy=True)
+        return (2.0 * self._cur_state - self._prev_state).astype(np.float32)
+
+
+if __name__ == "__main__":
+    model = PredictionModel()
+    print(f"Backend: {model.backend}")
+    for step in range(200):
+        dp = DataPoint(0, step, step >= 100, np.random.randn(32).astype(np.float32))
+        pred = model.predict(dp)
+        if pred is not None and step == 100:
+            print(f"step {step}: shape={pred.shape} dtype={pred.dtype} finite={np.all(np.isfinite(pred))}")
+    print("Local test passed!")

src/__init__.py

@@ -0,0 +1 @@
+# Wunder Fund RNN Challenge - Source Package

src/data/__init__.py

@@ -0,0 +1,30 @@
+"""Data package exports (lazy loading for the Torch-dependent dataset)."""
+
+import importlib
+
+from .protocol import create_dataset_manifest, create_sequence_folds, resolve_dataset_path, validate_wunder_dataframe
+from .causal_features import (
+    build_causal_tabular_features,
+    feature_schema_hash,
+    feature_schema_payload,
+    feature_width,
+)
+
+__all__ = [
+    'CausalStepDataset',
+    'create_causal_dataloaders',
+    'create_dataset_manifest',
+    'create_sequence_folds',
+    'resolve_dataset_path',
+    'validate_wunder_dataframe',
+    'build_causal_tabular_features',
+    'feature_schema_hash',
+    'feature_schema_payload',
+    'feature_width',
+]
+
+
+def __getattr__(name):
+    if name in {'CausalStepDataset', 'create_causal_dataloaders'}:
+        return getattr(importlib.import_module(".dataset", __name__), name)
+    raise AttributeError(f"module 'src.data' has no attribute {name!r}")

src/data/causal_features.py

@@ -0,0 +1,117 @@
+"""Causal online feature builders shared by Phase 2 tabular baselines."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+
+import numpy as np
+
+
+COMPACT_V1_WINDOWS = (3, 5, 10, 20, 50, 100)
+EXPANDED_V2_LAGS = (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610)
+EXPANDED_V2_EWMA_ALPHAS = (0.02, 0.05, 0.10, 0.20, 0.40)
+EXPANDED_V2_WINDOWS = (3, 5, 10, 20, 50, 100, 200, 500)
+
+
+def _validate_history(history: np.ndarray) -> np.ndarray:
+    history = np.asarray(history, dtype=np.float32)
+    if history.ndim != 2 or history.shape[1] != 32:
+        raise ValueError(f"Expected history shape (t, 32), got {history.shape}")
+    if len(history) == 0:
+        raise ValueError("history must contain at least one state")
+    return history
+
+
+def _ewma(history: np.ndarray, alpha: float) -> np.ndarray:
+    state = history[0].astype(np.float32, copy=True)
+    for row in history[1:]:
+        state = alpha * row + (1.0 - alpha) * state
+    return state.astype(np.float32)
+
+
+def build_compact_v1_features(history: np.ndarray) -> np.ndarray:
+    """Build the original compact raw-32 causal feature vector."""
+    history = _validate_history(history)
+    current = history[-1]
+    prev = history[-2] if len(history) >= 2 else np.zeros_like(current)
+    delta1 = current - prev
+
+    features = [current, delta1]
+    for window in COMPACT_V1_WINDOWS:
+        chunk = history[-window:]
+        mean = chunk.mean(axis=0)
+        features.extend([mean, chunk.std(axis=0), current - mean])
+    return np.concatenate(features).astype(np.float32)
+
+
+def build_expanded_v2_features(history: np.ndarray) -> np.ndarray:
+    """Build the Phase 2 max-safe expanded causal feature vector."""
+    history = _validate_history(history)
+    current = history[-1]
+    features = [current]
+
+    for lag in EXPANDED_V2_LAGS:
+        lagged = history[-1 - lag] if len(history) > lag else history[0]
+        features.extend([lagged, current - lagged])
+
+    for alpha in EXPANDED_V2_EWMA_ALPHAS:
+        smoothed = _ewma(history, alpha)
+        features.extend([smoothed, current - smoothed])
+
+    for window in EXPANDED_V2_WINDOWS:
+        chunk = history[-window:]
+        mean = chunk.mean(axis=0)
+        features.extend(
+            [
+                mean,
+                chunk.std(axis=0),
+                chunk.min(axis=0),
+                chunk.max(axis=0),
+                current - mean,
+            ]
+        )
+    return np.concatenate(features).astype(np.float32)
+
+
+def build_causal_tabular_features(history: np.ndarray, schema: str = "compact_v1") -> np.ndarray:
+    if schema == "compact_v1":
+        return build_compact_v1_features(history)
+    if schema == "expanded_v2":
+        return build_expanded_v2_features(history)
+    raise ValueError(f"Unknown feature schema: {schema}")
+
+
+def feature_width(schema: str) -> int:
+    return int(build_causal_tabular_features(np.zeros((1, 32), dtype=np.float32), schema).shape[0])
+
+
+def infer_feature_schema_from_width(width: int) -> str:
+    width = int(width)
+    for schema in ("compact_v1", "expanded_v2"):
+        if feature_width(schema) == width:
+            return schema
+    raise ValueError(f"No known causal feature schema has width {width}")
+
+
+def feature_schema_payload(schema: str) -> dict:
+    if schema == "compact_v1":
+        return {
+            "schema": schema,
+            "windows": list(COMPACT_V1_WINDOWS),
+            "width": feature_width(schema),
+        }
+    if schema == "expanded_v2":
+        return {
+            "schema": schema,
+            "lags": list(EXPANDED_V2_LAGS),
+            "ewma_alphas": list(EXPANDED_V2_EWMA_ALPHAS),
+            "windows": list(EXPANDED_V2_WINDOWS),
+            "width": feature_width(schema),
+        }
+    raise ValueError(f"Unknown feature schema: {schema}")
+
+
+def feature_schema_hash(schema: str) -> str:
+    payload = json.dumps(feature_schema_payload(schema), sort_keys=True).encode("utf-8")
+    return hashlib.sha256(payload).hexdigest()

src/data/dataset.py

@@ -0,0 +1,119 @@
+"""
+Causal step dataset for the Wunder Fund RNN Challenge.
+
+Provides CausalStepDataset (left-padded lookback windows for the official
+stepwise task) and a deterministic train/val loader factory. Targets are
+aligned so that the example at step t predicts state t+1, and only scored steps
+(need_prediction=True, with an available next state) are emitted.
+"""
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+from typing import Optional, List, Tuple
+
+from src.data.protocol import get_feature_columns, load_wunder_dataframe
+from src.utils.reproducibility import make_torch_generator, seed_worker
+
+
+class CausalStepDataset(Dataset):
+    """
+    Causal online dataset for the official stepwise task.
+
+    Each sample contains only states up to and including step t and predicts
+    state t+1. Samples are emitted only where need_prediction=True and t+1 is
+    available in the local parquet file. The hidden competition may request a
+    final step prediction whose target is unavailable locally; this dataset
+    deliberately excludes that row from training and local validation.
+    """
+
+    def __init__(
+        self,
+        parquet_path: str,
+        seq_ids: Optional[List[int]] = None,
+        lookback: int = 128,
+    ):
+        self.lookback = int(lookback)
+        if self.lookback <= 0:
+            raise ValueError("lookback must be positive")
+
+        df = load_wunder_dataframe(parquet_path, seq_ids=seq_ids)
+        self.feature_cols = get_feature_columns(df)
+        self.n_features = len(self.feature_cols)
+        self.samples = []
+
+        for seq_ix, seq_df in df.groupby("seq_ix", sort=True):
+            states = seq_df[self.feature_cols].to_numpy(dtype=np.float32)
+            need_prediction = seq_df["need_prediction"].to_numpy(dtype=bool)
+            steps = seq_df["step_in_seq"].to_numpy(dtype=np.int64)
+            for pos, step in enumerate(steps):
+                if not need_prediction[pos]:
+                    continue
+                target_pos = pos + 1
+                if target_pos >= len(states):
+                    continue
+                self.samples.append(
+                    {
+                        "seq_ix": int(seq_ix),
+                        "step_in_seq": int(step),
+                        "states": states,
+                        "position": int(pos),
+                        "target": states[target_pos],
+                    }
+                )
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> dict:
+        sample = self.samples[idx]
+        states = sample["states"]
+        pos = sample["position"]
+        start = max(0, pos + 1 - self.lookback)
+        history = states[start : pos + 1]
+
+        window = np.zeros((self.lookback, self.n_features), dtype=np.float32)
+        mask = np.zeros((self.lookback,), dtype=np.float32)
+        window[-len(history) :] = history
+        mask[-len(history) :] = 1.0
+
+        return {
+            "seq_ix": torch.tensor(sample["seq_ix"], dtype=torch.long),
+            "step_in_seq": torch.tensor(sample["step_in_seq"], dtype=torch.long),
+            "history": torch.tensor(window, dtype=torch.float32),
+            "history_mask": torch.tensor(mask, dtype=torch.float32),
+            "target": torch.tensor(sample["target"].astype(np.float32, copy=False), dtype=torch.float32),
+        }
+
+
+def create_causal_dataloaders(
+    parquet_path: str,
+    train_seq_ids: list[int],
+    val_seq_ids: list[int],
+    *,
+    lookback: int = 128,
+    batch_size: int = 256,
+    num_workers: int = 0,
+    seed: int = 42,
+) -> Tuple[DataLoader, DataLoader]:
+    """Create deterministic train/validation loaders for causal step examples."""
+    train_dataset = CausalStepDataset(parquet_path, seq_ids=train_seq_ids, lookback=lookback)
+    val_dataset = CausalStepDataset(parquet_path, seq_ids=val_seq_ids, lookback=lookback)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=False,
+        worker_init_fn=seed_worker,
+        generator=make_torch_generator(seed),
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=False,
+    )
+    return train_loader, val_loader

src/data/protocol.py

@@ -0,0 +1,229 @@
+"""Official-protocol checks, dataset manifests, and locked sequence folds."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import subprocess
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterable, Optional
+
+import numpy as np
+import pandas as pd
+
+
+REQUIRED_META_COLUMNS = ["seq_ix", "step_in_seq", "need_prediction"]
+EXPECTED_N_FEATURES = 32
+EXPECTED_SEQUENCE_LENGTH = 1000
+
+
+@dataclass(frozen=True)
+class DatasetManifest:
+    path: str
+    sha256: str
+    size_bytes: int
+    row_count: int
+    sequence_count: int
+    sequence_length: int
+    feature_columns: list[str]
+    need_prediction_counts_by_step: dict[str, int]
+    seq_ix_min: int
+    seq_ix_max: int
+
+
+@dataclass(frozen=True)
+class FoldManifest:
+    seed: int
+    n_folds: int
+    final_holdout_fold: int
+    folds: dict[str, list[int]]
+    train_dev_seq_ids: list[int]
+    final_holdout_seq_ids: list[int]
+    dataset_sha256: Optional[str] = None
+    git_commit: Optional[str] = None
+
+
+def resolve_dataset_path(preferred: str | Path = "data/raw/train.parquet") -> Path:
+    """Find the canonical raw train parquet without mutating the workspace."""
+    candidates = [
+        Path(preferred),
+        Path("data/raw/train.parquet"),
+        Path("data/train.parquet"),
+        Path("competition_package/datasets/train.parquet"),
+        Path("wnn_starterpack/competition_package/datasets/train.parquet"),
+        Path("datasets/train.parquet"),
+        Path("train.parquet"),
+    ]
+    for candidate in candidates:
+        if candidate.exists():
+            return candidate
+    raise FileNotFoundError(
+        "No train parquet found. Expected data/raw/train.parquet, "
+        "data/train.parquet, or competition_package/datasets/train.parquet."
+    )
+
+
+def sha256_file(path: str | Path, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def get_feature_columns(df: pd.DataFrame) -> list[str]:
+    """Return feature columns in scorer/parquet order, never lexicographic sort."""
+    return list(df.columns[len(REQUIRED_META_COLUMNS) :])
+
+
+def validate_wunder_dataframe(
+    df: pd.DataFrame,
+    *,
+    expected_n_features: int = EXPECTED_N_FEATURES,
+    expected_sequence_length: int = EXPECTED_SEQUENCE_LENGTH,
+) -> list[str]:
+    """Validate the official data contract and return feature columns."""
+    if list(df.columns[:3]) != REQUIRED_META_COLUMNS:
+        raise ValueError(
+            "First columns must be seq_ix, step_in_seq, need_prediction; "
+            f"got {list(df.columns[:3])}"
+        )
+
+    feature_cols = get_feature_columns(df)
+    if len(feature_cols) != expected_n_features:
+        raise ValueError(
+            f"Expected {expected_n_features} feature columns, got {len(feature_cols)}"
+        )
+    if df[REQUIRED_META_COLUMNS].isna().any().any():
+        raise ValueError("Metadata columns contain missing values")
+
+    lengths = df.groupby("seq_ix", sort=True).size()
+    bad_lengths = lengths[lengths != expected_sequence_length]
+    if not bad_lengths.empty:
+        raise ValueError(
+            "Every sequence must have exactly "
+            f"{expected_sequence_length} rows; bad seq_ix values: "
+            f"{bad_lengths.index.tolist()[:10]}"
+        )
+
+    expected_steps = np.arange(expected_sequence_length)
+    for seq_ix, seq_df in df.groupby("seq_ix", sort=False):
+        steps = seq_df["step_in_seq"].to_numpy()
+        if not np.array_equal(steps, expected_steps):
+            raise ValueError(
+                f"Sequence {seq_ix} must be ordered with steps 0.."
+                f"{expected_sequence_length - 1}"
+            )
+
+    if df[feature_cols].isna().any().any():
+        raise ValueError("Feature columns contain missing values")
+    return feature_cols
+
+
+def load_wunder_dataframe(path: str | Path, seq_ids: Optional[Iterable[int]] = None) -> pd.DataFrame:
+    """Read and validate the Wunder parquet file."""
+    df = pd.read_parquet(path)
+    if seq_ids is not None:
+        seq_set = set(int(seq_id) for seq_id in seq_ids)
+        df = df[df["seq_ix"].isin(seq_set)].copy()
+    df = df.sort_values(["seq_ix", "step_in_seq"]).reset_index(drop=True)
+    validate_wunder_dataframe(df)
+    return df
+
+
+def create_dataset_manifest(path: str | Path) -> DatasetManifest:
+    """Build a reproducibility manifest for a validated parquet file."""
+    path = Path(path)
+    df = load_wunder_dataframe(path)
+    feature_cols = get_feature_columns(df)
+    need_counts = (
+        df.groupby("step_in_seq")["need_prediction"]
+        .sum()
+        .astype(int)
+        .to_dict()
+    )
+    seq_ids = sorted(int(x) for x in df["seq_ix"].unique())
+    return DatasetManifest(
+        path=str(path),
+        sha256=sha256_file(path),
+        size_bytes=path.stat().st_size,
+        row_count=int(len(df)),
+        sequence_count=len(seq_ids),
+        sequence_length=int(df.groupby("seq_ix").size().iloc[0]),
+        feature_columns=feature_cols,
+        need_prediction_counts_by_step={str(k): int(v) for k, v in need_counts.items()},
+        seq_ix_min=min(seq_ids),
+        seq_ix_max=max(seq_ids),
+    )
+
+
+def write_json(path: str | Path, payload: object) -> None:
+    """Write dataclasses/dicts as stable JSON."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if hasattr(payload, "__dataclass_fields__"):
+        payload = asdict(payload)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def read_json(path: str | Path) -> dict:
+    return json.loads(Path(path).read_text(encoding="utf-8"))
+
+
+def current_git_commit() -> Optional[str]:
+    try:
+        return subprocess.check_output(
+            ["git", "rev-parse", "HEAD"],
+            text=True,
+            stderr=subprocess.DEVNULL,
+        ).strip()
+    except Exception:
+        return None
+
+
+def create_sequence_folds(
+    seq_ids: Iterable[int],
+    *,
+    n_folds: int = 5,
+    seed: int = 42,
+    final_holdout_fold: int = 4,
+    dataset_sha256: Optional[str] = None,
+) -> FoldManifest:
+    """Create deterministic folds by seq_ix with one untouched final holdout fold."""
+    seq_ids = np.array(sorted(int(x) for x in seq_ids), dtype=int)
+    if len(seq_ids) < n_folds:
+        raise ValueError("Need at least as many sequences as folds")
+    if final_holdout_fold < 0 or final_holdout_fold >= n_folds:
+        raise ValueError("final_holdout_fold must be in [0, n_folds)")
+
+    rng = np.random.RandomState(seed)
+    shuffled = seq_ids.copy()
+    rng.shuffle(shuffled)
+    split_ids = np.array_split(shuffled, n_folds)
+    folds = {
+        str(i): sorted(int(x) for x in split_ids[i].tolist())
+        for i in range(n_folds)
+    }
+    final_holdout = folds[str(final_holdout_fold)]
+    train_dev = sorted(
+        int(seq_id)
+        for fold_id, fold_seq_ids in folds.items()
+        if int(fold_id) != final_holdout_fold
+        for seq_id in fold_seq_ids
+    )
+    return FoldManifest(
+        seed=seed,
+        n_folds=n_folds,
+        final_holdout_fold=final_holdout_fold,
+        folds=folds,
+        train_dev_seq_ids=train_dev,
+        final_holdout_seq_ids=final_holdout,
+        dataset_sha256=dataset_sha256,
+        git_commit=current_git_commit(),
+    )
+
+
+def load_sequence_folds(path: str | Path) -> FoldManifest:
+    payload = read_json(path)
+    return FoldManifest(**payload)

src/evaluation/__init__.py

@@ -0,0 +1,3 @@
+from .stepwise import StepwiseScoreResult, StepwiseScorer
+
+__all__ = ["StepwiseScoreResult", "StepwiseScorer"]

src/evaluation/stepwise.py

@@ -0,0 +1,107 @@
+"""Official-style row-by-row evaluation without sequence-boundary leakage."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+from competition_package.utils import DataPoint
+from src.data.protocol import get_feature_columns, validate_wunder_dataframe
+from src.utils.metrics import compute_r2_per_feature, compute_r2_score
+
+
+@dataclass(frozen=True)
+class StepwiseScoreResult:
+    mean_r2: float
+    r2_per_feature: dict[str, float]
+    n_predictions_scored: int
+    n_predictions_requested: int
+    n_predictions_dropped_final_step: int
+    feature_columns: list[str]
+
+
+class StepwiseScorer:
+    """Replay held-out sequences through the competition `predict` API."""
+
+    def __init__(self, dataframe: pd.DataFrame):
+        self.dataset = dataframe.sort_values(["seq_ix", "step_in_seq"]).reset_index(drop=True)
+        self.feature_columns = validate_wunder_dataframe(self.dataset)
+        self.dim = len(self.feature_columns)
+
+    @classmethod
+    def from_parquet(cls, path: str, seq_ids: Optional[list[int]] = None) -> "StepwiseScorer":
+        df = pd.read_parquet(path)
+        if seq_ids is not None:
+            df = df[df["seq_ix"].isin(set(seq_ids))].copy()
+        return cls(df)
+
+    def score(self, model) -> StepwiseScoreResult:
+        predictions = []
+        targets = []
+        pending_prediction = None
+        pending_seq_ix = None
+        requested = 0
+        dropped_final = 0
+
+        for _, seq_df in self.dataset.groupby("seq_ix", sort=True):
+            pending_prediction = None
+            pending_seq_ix = None
+
+            seq_ix_values = seq_df["seq_ix"].to_numpy()
+            step_values = seq_df["step_in_seq"].to_numpy()
+            need_values = seq_df["need_prediction"].to_numpy()
+            state_values = seq_df[self.feature_columns].to_numpy(dtype=np.float32)
+
+            for pos in range(len(seq_df)):
+                seq_ix = int(seq_ix_values[pos])
+                step_in_seq = int(step_values[pos])
+                need_prediction = bool(need_values[pos])
+                state = state_values[pos]
+
+                if pending_prediction is not None:
+                    if pending_seq_ix != seq_ix:
+                        raise AssertionError("Internal scorer leaked across sequences")
+                    predictions.append(pending_prediction)
+                    targets.append(state)
+
+                data_point = DataPoint(seq_ix, step_in_seq, need_prediction, state)
+                next_prediction = model.predict(data_point)
+                self._check_prediction(data_point, next_prediction)
+                if need_prediction:
+                    requested += 1
+                pending_prediction = next_prediction
+                pending_seq_ix = seq_ix
+
+            if pending_prediction is not None:
+                dropped_final += 1
+
+        if not predictions:
+            raise ValueError("No predictions were scored")
+
+        y_pred = np.asarray(predictions, dtype=np.float64)
+        y_true = np.asarray(targets, dtype=np.float64)
+        return StepwiseScoreResult(
+            mean_r2=compute_r2_score(y_true, y_pred),
+            r2_per_feature=compute_r2_per_feature(y_true, y_pred, self.feature_columns),
+            n_predictions_scored=len(predictions),
+            n_predictions_requested=requested,
+            n_predictions_dropped_final_step=dropped_final,
+            feature_columns=self.feature_columns,
+        )
+
+    def _check_prediction(self, data_point: DataPoint, prediction: Optional[np.ndarray]) -> None:
+        if not data_point.need_prediction:
+            if prediction is not None:
+                raise ValueError(f"Prediction is not needed for {data_point}")
+            return
+        if prediction is None:
+            raise ValueError(f"Prediction is required for {data_point}")
+        if not isinstance(prediction, np.ndarray):
+            raise ValueError(f"Prediction must be np.ndarray, got {type(prediction)}")
+        if prediction.shape != (self.dim,):
+            raise ValueError(f"Prediction has wrong shape: {prediction.shape} != {(self.dim,)}")
+        if not np.isfinite(prediction).all():
+            raise ValueError("Prediction contains NaN or Inf")

src/models/__init__.py

@@ -0,0 +1,32 @@
+"""Model package exports (lazy to keep torch optional for tabular-only use)."""
+
+import importlib
+
+__all__ = [
+    'PersistencePredictionModel',
+    'MomentumPredictionModel',
+    'EWMAPredictionModel',
+    'TabularStatefulPredictionModel',
+    'build_causal_tabular_features',
+    'CausalGRUForecaster',
+    'CausalTCN',
+    'EnsemblePredictionModel',
+]
+
+
+def __getattr__(name):
+    if name in {
+        'PersistencePredictionModel',
+        'MomentumPredictionModel',
+        'EWMAPredictionModel',
+        'TabularStatefulPredictionModel',
+        'build_causal_tabular_features',
+    }:
+        return getattr(importlib.import_module(".baselines", __name__), name)
+    if name == 'CausalGRUForecaster':
+        return getattr(importlib.import_module(".sequence_models", __name__), name)
+    if name == 'CausalTCN':
+        return getattr(importlib.import_module(".tcn", __name__), name)
+    if name == 'EnsemblePredictionModel':
+        return getattr(importlib.import_module(".ensemble_predictor", __name__), name)
+    raise AttributeError(f"module 'src.models' has no attribute {name!r}")

src/models/baselines.py

@@ -0,0 +1,103 @@
+"""Causal baseline predictors for the Wunder stepwise API."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+from src.data.causal_features import build_causal_tabular_features
+
+
+class PersistencePredictionModel:
+    """Predict the current state as the next state."""
+
+    def __init__(self, n_features: int = 32):
+        self.n_features = n_features
+        self.current_seq: Optional[int] = None
+        self.last_state: Optional[np.ndarray] = None
+
+    def predict(self, data_point):
+        if self.current_seq != data_point.seq_ix:
+            self.current_seq = data_point.seq_ix
+            self.last_state = None
+        self.last_state = data_point.state.astype(np.float32, copy=True)
+        if not data_point.need_prediction:
+            return None
+        return self.last_state.astype(np.float32, copy=True)
+
+
+class MomentumPredictionModel:
+    """Predict next state with a first-difference momentum term."""
+
+    def __init__(self, alpha: float = 1.0, n_features: int = 32):
+        self.alpha = float(alpha)
+        self.n_features = n_features
+        self.current_seq: Optional[int] = None
+        self.previous_state: Optional[np.ndarray] = None
+        self.current_state: Optional[np.ndarray] = None
+
+    def predict(self, data_point):
+        if self.current_seq != data_point.seq_ix:
+            self.current_seq = data_point.seq_ix
+            self.previous_state = None
+            self.current_state = None
+
+        self.previous_state = self.current_state
+        self.current_state = data_point.state.astype(np.float32, copy=True)
+        if not data_point.need_prediction:
+            return None
+        if self.previous_state is None:
+            return self.current_state.astype(np.float32, copy=True)
+        pred = self.current_state + self.alpha * (self.current_state - self.previous_state)
+        return pred.astype(np.float32)
+
+
+class EWMAPredictionModel:
+    """Predict with an exponentially weighted moving average state."""
+
+    def __init__(self, alpha: float = 0.2, n_features: int = 32):
+        if not 0.0 < alpha <= 1.0:
+            raise ValueError("alpha must be in (0, 1]")
+        self.alpha = float(alpha)
+        self.n_features = n_features
+        self.current_seq: Optional[int] = None
+        self.ewma_state: Optional[np.ndarray] = None
+
+    def predict(self, data_point):
+        if self.current_seq != data_point.seq_ix:
+            self.current_seq = data_point.seq_ix
+            self.ewma_state = None
+        state = data_point.state.astype(np.float32, copy=False)
+        if self.ewma_state is None:
+            self.ewma_state = state.copy()
+        else:
+            self.ewma_state = self.alpha * state + (1.0 - self.alpha) * self.ewma_state
+        if not data_point.need_prediction:
+            return None
+        return self.ewma_state.astype(np.float32, copy=True)
+
+
+class TabularStatefulPredictionModel:
+    """Stateful wrapper for sklearn-style multi-output tabular regressors."""
+
+    def __init__(self, estimator, n_features: int = 32, feature_schema: str = "compact_v1"):
+        self.estimator = estimator
+        self.n_features = n_features
+        self.feature_schema = feature_schema
+        self.current_seq: Optional[int] = None
+        self.history: list[np.ndarray] = []
+
+    def predict(self, data_point):
+        if self.current_seq != data_point.seq_ix:
+            self.current_seq = data_point.seq_ix
+            self.history = []
+        self.history.append(data_point.state.astype(np.float32, copy=True))
+        if not data_point.need_prediction:
+            return None
+        features = build_causal_tabular_features(
+            np.asarray(self.history, dtype=np.float32),
+            schema=self.feature_schema,
+        )
+        pred = self.estimator.predict(features.reshape(1, -1))[0]
+        return np.asarray(pred, dtype=np.float32)

src/models/ensemble_predictor.py

@@ -0,0 +1,49 @@
+"""Uniform ensemble PredictionModel for the Wunder stepwise task.
+
+Both GRUStatefulPredictionModel and TabularStatefulPredictionModel expose the
+same contract: ``predict(data_point)`` advances per-sequence internal state and
+returns None during warm-up or an (n_features,) array when a prediction is
+needed. An ensemble is therefore just a list of member predictors plus an
+optional per-feature weight matrix. The same class powers local holdout
+evaluation and the packaged solution.py.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+
+class EnsemblePredictionModel:
+    def __init__(self, members: list, weights: Optional[np.ndarray] = None, n_features: int = 32):
+        """Args:
+            members: objects exposing predict(data_point) -> Optional[np.ndarray].
+            weights: optional (n_features, n_members) blend weights. If None, the
+                     members are averaged equally.
+        """
+        if not members:
+            raise ValueError("EnsemblePredictionModel needs at least one member")
+        self.members = members
+        self.n_features = n_features
+        self.weights = None
+        if weights is not None:
+            weights = np.asarray(weights, dtype=np.float32)
+            if weights.shape != (n_features, len(members)):
+                raise ValueError(
+                    f"weights shape {weights.shape} != {(n_features, len(members))}"
+                )
+            self.weights = weights
+
+    def predict(self, data_point) -> Optional[np.ndarray]:
+        # Always advance every member's state, even during warm-up.
+        member_preds = [m.predict(data_point) for m in self.members]
+        if not data_point.need_prediction:
+            return None
+        stacked = np.stack([np.asarray(p, dtype=np.float32) for p in member_preds], axis=0)  # (M, F)
+        if self.weights is not None:
+            # out[f] = sum_m stacked[m, f] * weights[f, m]
+            out = np.einsum("mf,fm->f", stacked, self.weights)
+        else:
+            out = stacked.mean(axis=0)
+        return out.astype(np.float32)

src/models/sequence_inference.py

@@ -0,0 +1,73 @@
+"""Online stateful inference wrapper for the causal GRU.
+
+Carries the GRU hidden state across steps so each prediction is O(1) work,
+which is what makes the model fit the 1-core / 60-minute submission budget.
+Used by both local stepwise evaluation and the packaged solution.py.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+
+from src.models.sequence_models import CausalGRUForecaster
+
+
+class GRUStatefulPredictionModel:
+    """Wrap a CausalGRUForecaster as a competition PredictionModel."""
+
+    def __init__(self, model: CausalGRUForecaster, n_features: int = 32):
+        self.model = model.eval()
+        self.n_features = n_features
+        self.current_seq: Optional[int] = None
+        self.h: Optional[torch.Tensor] = None
+
+    def reset(self):
+        self.current_seq = None
+        self.h = None
+
+    @torch.no_grad()
+    def predict(self, data_point):
+        if self.current_seq != data_point.seq_ix:
+            self.current_seq = data_point.seq_ix
+            self.h = None
+        state = torch.from_numpy(np.asarray(data_point.state, dtype=np.float32))
+        pred, self.h = self.model.step(state, self.h)
+        if not data_point.need_prediction:
+            return None
+        return pred.detach().numpy().astype(np.float32)
+
+
+def build_model_from_config(params: dict, model_type: str = "CausalGRUForecaster"):
+    if model_type == "CausalTCN":
+        from src.models.tcn import CausalTCN
+        return CausalTCN(
+            n_features=params.get("n_features", 32),
+            d_model=params.get("d_model", 192),
+            n_layers=params.get("n_layers", 6),
+            kernel_size=params.get("kernel_size", 3),
+            dropout=params.get("dropout", 0.1),
+            head_hidden=params.get("head_hidden"),
+        )
+    return CausalGRUForecaster(
+        n_features=params.get("n_features", 32),
+        d_model=params.get("d_model", 256),
+        n_layers=params.get("n_layers", 2),
+        dropout=params.get("dropout", 0.1),
+        head_hidden=params.get("head_hidden"),
+        rnn_type=params.get("rnn_type", "gru"),
+    )
+
+
+def load_gru_checkpoint(path: str | Path, map_location: str = "cpu"):
+    """Reconstruct a forecaster (GRU/LSTM/TCN) from a saved checkpoint."""
+    ckpt = torch.load(str(path), map_location=map_location, weights_only=False)
+    model_cfg = ckpt.get("config", {}).get("model", {})
+    params = model_cfg.get("params", {})
+    model = build_model_from_config(params, model_cfg.get("type", "CausalGRUForecaster"))
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+    return model

src/models/sequence_models.py

@@ -0,0 +1,127 @@
+"""Causal sequence models for the Wunder stepwise next-state task.
+
+Design constraints derived from the data analysis (see memory / report):
+  - Features are pre-whitened (~N(0,1), no heavy tails) so no internal
+    normalization is required; the model predicts the next *level* directly.
+  - Inference is online, one row at a time, on a single CPU core within a
+    60-minute budget. A unidirectional GRU carries hidden state across steps,
+    giving O(1) work per prediction.
+  - Training uses full 1000-step sequences with back-prop-through-time and a
+    loss masked to the scored steps (100..998). Because a GRU is a pure
+    recurrence, the full-sequence forward and the stepwise stateful forward
+    produce identical outputs — so what we train is exactly what we score.
+
+`forward(x, h0)` returns per-step predictions of the *next* state:
+    preds[:, t] predicts x[:, t + 1].
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class CausalGRUForecaster(nn.Module):
+    """Unidirectional GRU that predicts the next state at every step.
+
+    Args:
+        n_features: input/output width (32).
+        d_model: GRU hidden size and embedding width.
+        n_layers: stacked GRU layers.
+        dropout: dropout between GRU layers and inside the head.
+        head_hidden: width of the MLP head hidden layer (defaults to d_model).
+    """
+
+    def __init__(
+        self,
+        n_features: int = 32,
+        d_model: int = 256,
+        n_layers: int = 2,
+        dropout: float = 0.1,
+        head_hidden: Optional[int] = None,
+        rnn_type: str = "gru",
+    ):
+        super().__init__()
+        self.n_features = int(n_features)
+        self.d_model = int(d_model)
+        self.n_layers = int(n_layers)
+        self.rnn_type = str(rnn_type).lower()
+        head_hidden = int(head_hidden) if head_hidden else self.d_model
+
+        # Per-step input embedding (applied identically in batch and stepwise modes).
+        self.input_proj = nn.Linear(n_features, d_model)
+        self.input_norm = nn.LayerNorm(d_model)
+
+        rnn_cls = nn.LSTM if self.rnn_type == "lstm" else nn.GRU
+        self.gru = rnn_cls(
+            input_size=d_model,
+            hidden_size=d_model,
+            num_layers=n_layers,
+            batch_first=True,
+            dropout=dropout if n_layers > 1 else 0.0,
+        )
+
+        self.head = nn.Sequential(
+            nn.LayerNorm(d_model),
+            nn.Linear(d_model, head_hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(head_hidden, n_features),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        h0: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Args:
+            x: (batch, seq_len, n_features)
+            h0: optional initial hidden state (n_layers, batch, d_model)
+        Returns:
+            preds: (batch, seq_len, n_features) where preds[:, t] predicts x[:, t+1]
+            h_n:   final hidden state (n_layers, batch, d_model)
+        """
+        z = self.input_norm(self.input_proj(x))
+        out, h_n = self.gru(z, h0)
+        preds = self.head(out)
+        return preds, h_n
+
+    @torch.no_grad()
+    def step(
+        self,
+        state: torch.Tensor,
+        h: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Single online step.
+
+        Args:
+            state: (n_features,) or (1, n_features) current state.
+            h: previous hidden state (n_layers, 1, d_model) or None.
+        Returns:
+            pred: (n_features,) prediction of the next state.
+            h_new: updated hidden state.
+        """
+        if state.dim() == 1:
+            state = state.view(1, 1, -1)
+        elif state.dim() == 2:
+            state = state.unsqueeze(1)
+        preds, h_new = self.forward(state, h)
+        return preds[0, -1], h_new
+
+
+def build_gru(config: dict) -> CausalGRUForecaster:
+    """Construct a CausalGRUForecaster from a plain dict config."""
+    return CausalGRUForecaster(
+        n_features=config.get("n_features", 32),
+        d_model=config.get("d_model", 256),
+        n_layers=config.get("n_layers", 2),
+        dropout=config.get("dropout", 0.1),
+        head_hidden=config.get("head_hidden"),
+        rnn_type=config.get("rnn_type", "gru"),
+    )
+
+
+def count_parameters(model: nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)

src/models/tcn.py

@@ -0,0 +1,115 @@
+"""Causal Temporal Convolutional Network for next-state prediction.
+
+A different inductive bias from the RNNs (dilated causal convolutions instead of
+a recurrence), which decorrelates its errors from the GRU/LSTM members and so
+adds more in an ensemble than another RNN seed.
+
+Exposes the same interface as CausalGRUForecaster so it is a drop-in for the
+trainer, the stateful inference wrapper, and the ensembler:
+  forward(x, h0=None) -> (preds, None)   where preds[:, t] predicts x[:, t+1]
+  step(state, h)      -> (pred, buffer)  buffer holds the last receptive-field states
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class _ChannelLayerNorm(nn.Module):
+    """LayerNorm over the channel dim of a (B, C, T) tensor, independent per
+    timestep — causal, unlike GroupNorm/BatchNorm which mix statistics across T."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x):  # (B, C, T)
+        return self.norm(x.transpose(1, 2)).transpose(1, 2)
+
+
+class _CausalConvBlock(nn.Module):
+    def __init__(self, d_model: int, kernel_size: int, dilation: int, dropout: float):
+        super().__init__()
+        self.pad = (kernel_size - 1) * dilation
+        self.conv1 = nn.Conv1d(d_model, d_model, kernel_size, dilation=dilation, padding=self.pad)
+        self.conv2 = nn.Conv1d(d_model, d_model, kernel_size, dilation=dilation, padding=self.pad)
+        self.norm1 = _ChannelLayerNorm(d_model)
+        self.norm2 = _ChannelLayerNorm(d_model)
+        self.act = nn.GELU()
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x):  # x: (B, d, T)
+        T = x.shape[-1]
+        h = self.conv1(x)[..., :T]      # trim right (future) padding -> causal
+        h = self.drop(self.act(self.norm1(h)))
+        h = self.conv2(h)[..., :T]
+        h = self.drop(self.act(self.norm2(h)))
+        return x + h
+
+
+class CausalTCN(nn.Module):
+    def __init__(
+        self,
+        n_features: int = 32,
+        d_model: int = 192,
+        n_layers: int = 6,
+        kernel_size: int = 3,
+        dropout: float = 0.1,
+        head_hidden: Optional[int] = None,
+    ):
+        super().__init__()
+        self.n_features = int(n_features)
+        self.d_model = int(d_model)
+        self.n_layers = int(n_layers)
+        self.kernel_size = int(kernel_size)
+        head_hidden = int(head_hidden) if head_hidden else d_model
+
+        self.input_proj = nn.Linear(n_features, d_model)
+        self.input_norm = nn.LayerNorm(d_model)
+        self.blocks = nn.ModuleList([
+            _CausalConvBlock(d_model, kernel_size, dilation=2 ** i, dropout=dropout)
+            for i in range(n_layers)
+        ])
+        self.head = nn.Sequential(
+            nn.LayerNorm(d_model),
+            nn.Linear(d_model, head_hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(head_hidden, n_features),
+        )
+        # receptive field for buffered inference
+        self.receptive_field = 1 + 2 * sum((kernel_size - 1) * (2 ** i) for i in range(n_layers))
+
+    def forward(self, x: torch.Tensor, h0=None) -> Tuple[torch.Tensor, None]:
+        z = self.input_norm(self.input_proj(x))  # (B, T, d)
+        z = z.transpose(1, 2)                     # (B, d, T)
+        for block in self.blocks:
+            z = block(z)
+        z = z.transpose(1, 2)                     # (B, T, d)
+        return self.head(z), None
+
+    @torch.no_grad()
+    def step(self, state: torch.Tensor, h: Optional[torch.Tensor] = None):
+        """h is the running buffer of recent states (1, L, F)."""
+        if state.dim() == 1:
+            state = state.view(1, 1, -1)
+        elif state.dim() == 2:
+            state = state.unsqueeze(1)
+        buf = state if h is None else torch.cat([h, state], dim=1)
+        buf = buf[:, -self.receptive_field:]
+        preds, _ = self.forward(buf)
+        return preds[0, -1], buf
+
+
+def build_tcn(config: dict) -> CausalTCN:
+    return CausalTCN(
+        n_features=config.get("n_features", 32),
+        d_model=config.get("d_model", 192),
+        n_layers=config.get("n_layers", 6),
+        kernel_size=config.get("kernel_size", 3),
+        dropout=config.get("dropout", 0.1),
+        head_hidden=config.get("head_hidden"),
+    )

src/training/__init__.py

@@ -0,0 +1,18 @@
+"""Training package exports (lazy loading for the Torch-dependent trainer)."""
+
+import importlib
+
+__all__ = [
+    'build_tabular_supervised_matrix',
+    'create_tabular_estimator',
+    'TrainConfig',
+    'train_sequence_model',
+]
+
+
+def __getattr__(name):
+    if name in {'build_tabular_supervised_matrix', 'create_tabular_estimator'}:
+        return getattr(importlib.import_module(".baseline_training", __name__), name)
+    if name in {'TrainConfig', 'train_sequence_model'}:
+        return getattr(importlib.import_module(".sequence_trainer", __name__), name)
+    raise AttributeError(f"module 'src.training' has no attribute {name!r}")

src/training/baseline_training.py

@@ -0,0 +1,68 @@
+"""Training utilities for Phase 2 causal tabular baselines."""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+from src.data.protocol import get_feature_columns, load_wunder_dataframe
+from src.models.baselines import build_causal_tabular_features
+
+
+def build_tabular_supervised_matrix(
+    parquet_path: str,
+    seq_ids: list[int],
+    *,
+    feature_schema: str = "compact_v1",
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build causal tabular features and next-state targets for selected sequences."""
+    df = load_wunder_dataframe(parquet_path, seq_ids=seq_ids)
+    feature_cols = get_feature_columns(df)
+    rows = []
+    targets = []
+    for _, seq_df in df.groupby("seq_ix", sort=True):
+        states = seq_df[feature_cols].to_numpy(dtype=np.float32)
+        need_prediction = seq_df["need_prediction"].to_numpy(dtype=bool)
+        for pos in range(len(states) - 1):
+            if not need_prediction[pos]:
+                continue
+            rows.append(build_causal_tabular_features(states[: pos + 1], schema=feature_schema))
+            targets.append(states[pos + 1])
+    if not rows:
+        raise ValueError("No supervised rows were created")
+    return np.vstack(rows).astype(np.float32), np.vstack(targets).astype(np.float32)
+
+
+def create_tabular_estimator(model_type: str, seed: int = 42, n_jobs: int = 1):
+    """Create a deterministic sklearn-compatible multi-output regressor."""
+    model_type = model_type.lower()
+    if model_type == "ridge":
+        from sklearn.linear_model import Ridge
+
+        return Ridge(alpha=1.0, random_state=seed)
+    if model_type == "ridge_cv":
+        from sklearn.linear_model import RidgeCV
+        from sklearn.multioutput import MultiOutputRegressor
+
+        base = RidgeCV(alphas=np.logspace(-4, 4, 17))
+        return MultiOutputRegressor(base, n_jobs=max(1, int(n_jobs)))
+    if model_type == "elasticnet":
+        from sklearn.linear_model import MultiTaskElasticNet
+
+        return MultiTaskElasticNet(alpha=0.0005, l1_ratio=0.05, random_state=seed, max_iter=5000)
+    if model_type == "lightgbm":
+        import lightgbm as lgb
+        from sklearn.multioutput import MultiOutputRegressor
+
+        base = lgb.LGBMRegressor(
+            n_estimators=300,
+            learning_rate=0.03,
+            num_leaves=31,
+            subsample=0.9,
+            colsample_bytree=0.9,
+            random_state=seed,
+            n_jobs=1,
+            verbose=-1,
+        )
+        return MultiOutputRegressor(base, n_jobs=max(1, int(n_jobs)))
+    raise ValueError(f"Unknown tabular model type: {model_type}")

src/training/sequence_trainer.py

@@ -0,0 +1,199 @@
+"""Full-sequence BPTT trainer for the causal GRU next-state model.
+
+Trains on whole 1000-step sequences, computing a masked MSE loss only on the
+scored steps (need_prediction True, i.e. current steps 100..998 -> targets
+101..999). Validation R2 is computed with a batched full-sequence forward,
+which is numerically identical to the stepwise stateful replay because a GRU
+is a pure recurrence.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from src.data.protocol import get_feature_columns, load_wunder_dataframe
+from src.models.sequence_models import CausalGRUForecaster, count_parameters
+from src.models.tcn import CausalTCN
+from src.utils.metrics import compute_r2_per_feature, compute_r2_score
+from src.utils.reproducibility import set_global_seed
+
+
+def load_full_sequences(parquet_path: str, seq_ids: list[int]):
+    """Return (states, need, seq_ids) as dense arrays.
+
+    states: (N, 1000, 32) float32
+    need:   (N, 1000) bool   (need_prediction flag per step)
+    """
+    df = load_wunder_dataframe(parquet_path, seq_ids=seq_ids)
+    cols = get_feature_columns(df)
+    states_list, need_list, ids = [], [], []
+    for sid, g in df.sort_values(["seq_ix", "step_in_seq"]).groupby("seq_ix", sort=True):
+        states_list.append(g[cols].to_numpy(dtype=np.float32))
+        need_list.append(g["need_prediction"].to_numpy(dtype=bool))
+        ids.append(int(sid))
+    states = np.stack(states_list).astype(np.float32)
+    need = np.stack(need_list)
+    return states, need, ids
+
+
+@dataclass
+class TrainConfig:
+    d_model: int = 256
+    n_layers: int = 2
+    dropout: float = 0.1
+    head_hidden: Optional[int] = None
+    epochs: int = 40
+    batch_size: int = 32
+    lr: float = 1.0e-3
+    weight_decay: float = 1.0e-4
+    grad_clip: float = 1.0
+    warmup_frac: float = 0.1
+    seed: int = 42
+    threads: int = 8
+    patience: int = 12
+    n_features: int = 32
+    device: str = "cpu"
+    rnn_type: str = "gru"
+    arch: str = "rnn"        # "rnn" or "tcn"
+    kernel_size: int = 3     # tcn only
+
+
+def _masked_mse(preds: torch.Tensor, states: torch.Tensor, need: torch.Tensor) -> torch.Tensor:
+    # preds[:, t] predicts states[:, t+1]
+    pred = preds[:, :-1, :]
+    target = states[:, 1:, :]
+    mask = need[:, :-1]  # (B, T-1)
+    diff2 = (pred - target) ** 2  # (B, T-1, F)
+    m = mask.unsqueeze(-1).to(diff2.dtype)
+    return (diff2 * m).sum() / (m.sum() * preds.shape[-1])
+
+
+@torch.no_grad()
+def _eval_r2(model, states_t, need_t, feature_cols, batch_size=64):
+    model.eval()
+    preds_all, tgts_all = [], []
+    n = states_t.shape[0]
+    for i in range(0, n, batch_size):
+        sb = states_t[i : i + batch_size]
+        nb = need_t[i : i + batch_size]
+        preds, _ = model(sb)
+        pred = preds[:, :-1, :]
+        target = sb[:, 1:, :]
+        mask = nb[:, :-1]
+        preds_all.append(pred[mask].cpu().numpy())
+        tgts_all.append(target[mask].cpu().numpy())
+    y_pred = np.concatenate(preds_all).astype(np.float64)
+    y_true = np.concatenate(tgts_all).astype(np.float64)
+    mean_r2 = compute_r2_score(y_true, y_pred)
+    per_feat = compute_r2_per_feature(y_true, y_pred, feature_cols)
+    return mean_r2, per_feat, y_true, y_pred
+
+
+def train_sequence_model(
+    data_path: str,
+    train_ids: list[int],
+    val_ids: list[int],
+    cfg: TrainConfig,
+    feature_cols: Optional[list[str]] = None,
+    log_every: int = 1,
+    verbose: bool = True,
+):
+    set_global_seed(cfg.seed, deterministic_torch=True, seed_torch=True)
+    torch.set_num_threads(int(cfg.threads))
+    device = torch.device(cfg.device if (cfg.device != "cuda" or torch.cuda.is_available()) else "cpu")
+    if str(device) != cfg.device:
+        print(f"requested device '{cfg.device}' unavailable; using {device}")
+
+    s_tr, n_tr, _ = load_full_sequences(data_path, train_ids)
+    s_va, n_va, va_ids = load_full_sequences(data_path, val_ids)
+    if feature_cols is None:
+        feature_cols = [str(i) for i in range(cfg.n_features)]
+
+    states_tr = torch.from_numpy(s_tr).to(device)
+    need_tr = torch.from_numpy(n_tr).to(device)
+    states_va = torch.from_numpy(s_va).to(device)
+    need_va = torch.from_numpy(n_va).to(device)
+
+    if cfg.arch == "tcn":
+        model = CausalTCN(
+            n_features=cfg.n_features, d_model=cfg.d_model, n_layers=cfg.n_layers,
+            kernel_size=cfg.kernel_size, dropout=cfg.dropout, head_hidden=cfg.head_hidden,
+        ).to(device)
+    else:
+        model = CausalGRUForecaster(
+            n_features=cfg.n_features, d_model=cfg.d_model, n_layers=cfg.n_layers,
+            dropout=cfg.dropout, head_hidden=cfg.head_hidden, rnn_type=cfg.rnn_type,
+        ).to(device)
+    if verbose:
+        print(f"model params: {count_parameters(model):,}  train_seqs={len(train_ids)} val_seqs={len(val_ids)}")
+
+    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+    n_batches = max(1, (states_tr.shape[0] + cfg.batch_size - 1) // cfg.batch_size)
+    total_steps = cfg.epochs * n_batches
+    sched = torch.optim.lr_scheduler.OneCycleLR(
+        opt, max_lr=cfg.lr, total_steps=total_steps,
+        pct_start=cfg.warmup_frac, anneal_strategy="cos",
+    )
+
+    g = torch.Generator()
+    g.manual_seed(cfg.seed)
+    best_r2 = -1e9
+    best_state = None
+    best_per_feat = None
+    best_oof = None
+    history = []
+    bad_epochs = 0
+
+    for epoch in range(cfg.epochs):
+        model.train()
+        perm = torch.randperm(states_tr.shape[0], generator=g)
+        ep_loss = 0.0
+        t0 = time.perf_counter()
+        for bi in range(n_batches):
+            idx = perm[bi * cfg.batch_size : (bi + 1) * cfg.batch_size]
+            sb = states_tr[idx]
+            nb = need_tr[idx]
+            opt.zero_grad(set_to_none=True)
+            preds, _ = model(sb)
+            loss = _masked_mse(preds, sb, nb)
+            loss.backward()
+            nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
+            opt.step()
+            sched.step()
+            ep_loss += float(loss.detach())
+        ep_loss /= n_batches
+
+        val_r2, per_feat, y_true, y_pred = _eval_r2(model, states_va, need_va, feature_cols)
+        dt = time.perf_counter() - t0
+        history.append({"epoch": epoch, "train_loss": ep_loss, "val_r2": val_r2, "sec": dt})
+        if verbose and (epoch % log_every == 0):
+            print(f"epoch {epoch:3d}  loss={ep_loss:.5f}  val_R2={val_r2:.5f}  lr={sched.get_last_lr()[0]:.2e}  {dt:.1f}s")
+
+        if val_r2 > best_r2:
+            best_r2 = val_r2
+            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+            best_per_feat = per_feat
+            best_oof = (y_true, y_pred)
+            bad_epochs = 0
+        else:
+            bad_epochs += 1
+            if bad_epochs >= cfg.patience:
+                if verbose:
+                    print(f"early stop at epoch {epoch} (best val_R2={best_r2:.5f})")
+                break
+
+    return {
+        "best_val_r2": best_r2,
+        "best_per_feature": best_per_feat,
+        "best_state_dict": best_state,
+        "history": history,
+        "oof": best_oof,
+        "val_seq_ids": va_ids,
+        "feature_cols": feature_cols,
+    }

src/utils/__init__.py

@@ -0,0 +1,25 @@
+from .metrics import compute_manual_r2_score, compute_r2_per_feature, compute_r2_score
+from .reproducibility import make_torch_generator, seed_worker, set_global_seed
+from .hardware import (
+    HardwareProfile,
+    MaxSafeConfig,
+    ResourceLimitError,
+    assert_within_max_safe,
+    build_hardware_profile,
+    calibrate_batch_size,
+)
+
+__all__ = [
+    'compute_r2_score',
+    'compute_r2_per_feature',
+    'compute_manual_r2_score',
+    'set_global_seed',
+    'seed_worker',
+    'make_torch_generator',
+    'HardwareProfile',
+    'MaxSafeConfig',
+    'ResourceLimitError',
+    'build_hardware_profile',
+    'assert_within_max_safe',
+    'calibrate_batch_size',
+]

src/utils/hardware.py

@@ -0,0 +1,279 @@
+"""Hardware profiling and max-safe guard helpers for local-only experiments."""
+
+from __future__ import annotations
+
+import ctypes
+import contextlib
+import json
+import os
+import platform
+import shutil
+import subprocess
+import tempfile
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Callable, Iterable, Optional
+
+
+BYTES_PER_GB = 1024**3
+
+
+@dataclass(frozen=True)
+class MemoryInfo:
+    total_gb: float
+    available_gb: Optional[float]
+    used_fraction: Optional[float]
+
+
+@dataclass(frozen=True)
+class GpuInfo:
+    name: str
+    memory_total_mb: int
+    memory_used_mb: Optional[int]
+    temperature_c: Optional[int]
+    driver_version: Optional[str]
+
+
+@dataclass(frozen=True)
+class TorchInfo:
+    installed: bool
+    version: Optional[str]
+    cuda_available: bool
+    cuda_version: Optional[str]
+    device_name: Optional[str]
+    device_memory_gb: Optional[float]
+    error: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class MaxSafeConfig:
+    cpu_workers: int
+    dataloader_workers: int
+    max_dataloader_workers: int
+    ram_limit_fraction: float
+    ram_target_gb: float
+    gpu_vram_fraction: float
+    gpu_target_vram_gb: Optional[float]
+    gpu_pause_temp_c: int
+    gpu_abort_temp_c: int
+    mixed_precision: bool
+    gradient_accumulation: bool
+    batch_candidates: list[int]
+
+
+@dataclass(frozen=True)
+class HardwareProfile:
+    profile_name: str
+    platform: str
+    python: str
+    cpu_name: str
+    logical_cpu_count: int
+    memory: MemoryInfo
+    gpus: list[GpuInfo]
+    torch: TorchInfo
+    max_safe: MaxSafeConfig
+
+
+class ResourceLimitError(RuntimeError):
+    """Raised when a max-safe guard would abort a run."""
+
+
+def _round_gb(value_bytes: int) -> float:
+    return round(value_bytes / BYTES_PER_GB, 2)
+
+
+def get_memory_info() -> MemoryInfo:
+    """Return physical memory using stdlib-only platform APIs."""
+    if platform.system().lower() == "windows":
+        class MEMORYSTATUSEX(ctypes.Structure):
+            _fields_ = [
+                ("dwLength", ctypes.c_ulong),
+                ("dwMemoryLoad", ctypes.c_ulong),
+                ("ullTotalPhys", ctypes.c_ulonglong),
+                ("ullAvailPhys", ctypes.c_ulonglong),
+                ("ullTotalPageFile", ctypes.c_ulonglong),
+                ("ullAvailPageFile", ctypes.c_ulonglong),
+                ("ullTotalVirtual", ctypes.c_ulonglong),
+                ("ullAvailVirtual", ctypes.c_ulonglong),
+                ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
+            ]
+
+        status = MEMORYSTATUSEX()
+        status.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
+        ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status))
+        used_fraction = 1.0 - (status.ullAvailPhys / status.ullTotalPhys)
+        return MemoryInfo(
+            total_gb=_round_gb(status.ullTotalPhys),
+            available_gb=_round_gb(status.ullAvailPhys),
+            used_fraction=round(used_fraction, 4),
+        )
+
+    if hasattr(os, "sysconf") and "SC_PAGE_SIZE" in os.sysconf_names:
+        page_size = os.sysconf("SC_PAGE_SIZE")
+        pages = os.sysconf("SC_PHYS_PAGES")
+        return MemoryInfo(total_gb=_round_gb(page_size * pages), available_gb=None, used_fraction=None)
+
+    return MemoryInfo(total_gb=0.0, available_gb=None, used_fraction=None)
+
+
+def get_nvidia_gpus() -> list[GpuInfo]:
+    """Read NVIDIA GPU status through nvidia-smi when available."""
+    if not shutil.which("nvidia-smi"):
+        return []
+    cmd = [
+        "nvidia-smi",
+        "--query-gpu=name,memory.total,memory.used,temperature.gpu,driver_version",
+        "--format=csv,noheader,nounits",
+    ]
+    try:
+        raw = subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL, timeout=10)
+    except Exception:
+        return []
+
+    gpus = []
+    for line in raw.splitlines():
+        parts = [part.strip() for part in line.split(",")]
+        if len(parts) < 5:
+            continue
+        name, total_mb, used_mb, temp_c, driver = parts[:5]
+        gpus.append(
+            GpuInfo(
+                name=name,
+                memory_total_mb=int(float(total_mb)),
+                memory_used_mb=int(float(used_mb)),
+                temperature_c=int(float(temp_c)),
+                driver_version=driver,
+            )
+        )
+    return gpus
+
+
+def get_torch_info() -> TorchInfo:
+    """Return PyTorch/CUDA availability without making Torch mandatory."""
+    try:
+        with open(os.devnull, "w", encoding="utf-8") as devnull, contextlib.redirect_stderr(devnull):
+            import torch
+
+        cuda_available = bool(torch.cuda.is_available())
+        device_name = None
+        device_memory_gb = None
+        if cuda_available:
+            props = torch.cuda.get_device_properties(0)
+            device_name = torch.cuda.get_device_name(0)
+            device_memory_gb = round(props.total_memory / BYTES_PER_GB, 2)
+        return TorchInfo(
+            installed=True,
+            version=str(torch.__version__),
+            cuda_available=cuda_available,
+            cuda_version=getattr(torch.version, "cuda", None),
+            device_name=device_name,
+            device_memory_gb=device_memory_gb,
+        )
+    except Exception as exc:
+        return TorchInfo(
+            installed=False,
+            version=None,
+            cuda_available=False,
+            cuda_version=None,
+            device_name=None,
+            device_memory_gb=None,
+            error=repr(exc),
+        )
+
+
+def build_max_safe_config(memory: MemoryInfo, gpus: list[GpuInfo]) -> MaxSafeConfig:
+    """Choose conservative max-safe defaults for this laptop class."""
+    logical = os.cpu_count() or 1
+    cpu_workers = max(1, min(12, logical - 4 if logical > 4 else logical))
+    gpu_target_vram_gb = None
+    if gpus:
+        gpu_target_vram_gb = round((gpus[0].memory_total_mb / 1024) * 0.85, 2)
+    return MaxSafeConfig(
+        cpu_workers=cpu_workers,
+        dataloader_workers=2,
+        max_dataloader_workers=4,
+        ram_limit_fraction=0.90,
+        ram_target_gb=min(13.0, round(memory.total_gb * 0.82, 2)) if memory.total_gb else 13.0,
+        gpu_vram_fraction=0.85,
+        gpu_target_vram_gb=gpu_target_vram_gb,
+        gpu_pause_temp_c=82,
+        gpu_abort_temp_c=87,
+        mixed_precision=True,
+        gradient_accumulation=True,
+        batch_candidates=[512, 256, 128, 64, 32, 16, 8],
+    )
+
+
+def build_hardware_profile(profile_name: str = "max_safe") -> HardwareProfile:
+    memory = get_memory_info()
+    gpus = get_nvidia_gpus()
+    return HardwareProfile(
+        profile_name=profile_name,
+        platform=platform.platform(),
+        python=platform.python_version(),
+        cpu_name=platform.processor() or platform.machine(),
+        logical_cpu_count=os.cpu_count() or 1,
+        memory=memory,
+        gpus=gpus,
+        torch=get_torch_info(),
+        max_safe=build_max_safe_config(memory, gpus),
+    )
+
+
+def assert_within_max_safe(profile: HardwareProfile) -> None:
+    """Abort before a run if current memory/temperature already violates guardrails."""
+    if profile.memory.used_fraction is not None and profile.memory.used_fraction >= profile.max_safe.ram_limit_fraction:
+        raise ResourceLimitError(
+            f"System RAM usage {profile.memory.used_fraction:.1%} exceeds "
+            f"{profile.max_safe.ram_limit_fraction:.0%} guard"
+        )
+    if profile.gpus:
+        gpu = profile.gpus[0]
+        if gpu.temperature_c is not None and gpu.temperature_c >= profile.max_safe.gpu_abort_temp_c:
+            raise ResourceLimitError(
+                f"GPU temperature {gpu.temperature_c}C exceeds "
+                f"{profile.max_safe.gpu_abort_temp_c}C abort guard"
+            )
+
+
+def calibrate_batch_size(
+    trial_fn: Callable[[int], None],
+    candidates: Iterable[int],
+    *,
+    repeats: int = 2,
+) -> Optional[int]:
+    """
+    Choose the largest candidate batch that survives repeated trial calls.
+
+    `trial_fn` should raise RuntimeError for OOM-like failures. This helper is
+    intentionally framework-agnostic so tests can simulate OOM without Torch.
+    """
+    for batch_size in candidates:
+        try:
+            for _ in range(repeats):
+                trial_fn(int(batch_size))
+            return int(batch_size)
+        except RuntimeError:
+            continue
+    return None
+
+
+def atomic_write_json(path: str | Path, payload: object) -> None:
+    """Write JSON atomically so interrupted runs do not corrupt manifests."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if hasattr(payload, "__dataclass_fields__"):
+        payload = asdict(payload)
+    with tempfile.NamedTemporaryFile("w", delete=False, dir=str(path.parent), encoding="utf-8") as tmp:
+        json.dump(payload, tmp, indent=2, sort_keys=True)
+        tmp.write("\n")
+        tmp_path = Path(tmp.name)
+    tmp_path.replace(path)
+
+
+def dataclass_to_dict(value: object) -> dict:
+    if hasattr(value, "__dataclass_fields__"):
+        return asdict(value)
+    if isinstance(value, dict):
+        return value
+    raise TypeError(f"Unsupported JSON payload: {type(value)}")

src/utils/metrics.py

@@ -0,0 +1,59 @@
+"""R2 metrics for the Wunder Fund RNN Challenge.
+
+The authoritative local metric mirrors ``competition_package.utils``:
+``sklearn.metrics.r2_score`` is computed independently per feature and then
+averaged.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from sklearn.metrics import r2_score
+
+
+def _validate_shapes(y_true: np.ndarray, y_pred: np.ndarray) -> None:
+    if y_true.shape != y_pred.shape:
+        raise ValueError(
+            f"Shape mismatch: y_true {y_true.shape} vs y_pred {y_pred.shape}"
+        )
+    if y_true.ndim != 2:
+        raise ValueError(f"Expected 2D arrays, got y_true.ndim={y_true.ndim}")
+
+
+def compute_r2_score(y_true: np.ndarray, y_pred: np.ndarray) -> float:
+    """Compute mean per-feature R2 using sklearn, matching the starter pack."""
+    return float(np.mean(list(compute_r2_per_feature(y_true, y_pred).values())))
+
+
+def compute_r2_per_feature(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    feature_names: list[str] | None = None,
+) -> dict[str, float]:
+    """Compute sklearn R2 independently for each output feature."""
+    _validate_shapes(y_true, y_pred)
+    n_features = y_true.shape[1]
+    if feature_names is None:
+        feature_names = [f"feature_{i}" for i in range(n_features)]
+    if len(feature_names) != n_features:
+        raise ValueError("feature_names length must match prediction width")
+
+    return {
+        feature_names[i]: float(r2_score(y_true[:, i], y_pred[:, i]))
+        for i in range(n_features)
+    }
+
+
+def compute_manual_r2_score(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    epsilon: float = 1e-8,
+) -> float:
+    """Legacy epsilon-stabilized R2 helper for diagnostics only."""
+    _validate_shapes(y_true, y_pred)
+    scores = []
+    for i in range(y_true.shape[1]):
+        ss_res = np.sum((y_true[:, i] - y_pred[:, i]) ** 2)
+        ss_tot = np.sum((y_true[:, i] - y_true[:, i].mean()) ** 2)
+        scores.append(1 - (ss_res / (ss_tot + epsilon)))
+    return float(np.mean(scores))

src/utils/reproducibility.py

@@ -0,0 +1,56 @@
+"""Deterministic runtime helpers for Phase 2 experiments."""
+
+from __future__ import annotations
+
+import os
+import random
+from typing import Optional
+
+import numpy as np
+
+
+def set_global_seed(seed: int, deterministic_torch: bool = True, seed_torch: bool = True) -> None:
+    """Seed Python, NumPy, and PyTorch when available."""
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+    if not seed_torch:
+        return
+
+    try:
+        import torch
+
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+        if deterministic_torch:
+            torch.use_deterministic_algorithms(True, warn_only=True)
+            torch.backends.cudnn.benchmark = False
+            torch.backends.cudnn.deterministic = True
+    except ImportError:
+        return
+
+
+def seed_worker(worker_id: int) -> None:
+    """Seed DataLoader workers from PyTorch's worker seed."""
+    try:
+        import torch
+
+        worker_seed = torch.initial_seed() % 2**32
+    except ImportError:
+        worker_seed = worker_id
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+
+def make_torch_generator(seed: int) -> Optional[object]:
+    """Return a seeded torch.Generator, or None if torch is unavailable."""
+    try:
+        import torch
+
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        return generator
+    except ImportError:
+        return None