Fix the walk-forward split for a trading backtest

Review the AI's codeEasy10 minCitadel and Bridgewater Associates style

Your task

An AI validated a stock trading signal with KFold(shuffle=True) on rows sorted by date. Shuffling put future days into the training folds, so the backtest looked excellent and live trading lost money. The fix is a walk-forward split: train only on the past, test on the block after it, and leave an embargo gap between them, because rolling features near the boundary overlap.

An assistant wrote walk_forward_folds(n_rows, n_splits, embargo). It is in the editor. It has a bug. Find it and fix it.

Input and output:

  • n_rows: the number of rows, in date order, so a row's index is its place in time.
  • n_splits: how many test blocks to make.
  • embargo: how many rows to leave out between training and each test block.
  • Return one fold per test block, each [train_start, train_end, test_start, test_end]. Ends are exclusive: a range covers rows from its start up to, but not including, its end.

Rules the tests check:

  • Return an empty list if n_rows or n_splits is 0 or less, or if embargo is negative.
  • The block size is n_rows divided by n_splits + 1, rounded down. If it is 0, return an empty list.
  • The test blocks sit at the end of the data, in order. Fold f, counting from 0, starts its test block at n_rows - (n_splits - f) * size and ends it size rows later, so the last block ends at n_rows.
  • Training starts at 0 and ends at the test start minus embargo, but never below 0.

Examples

walk_forward_folds(n_rows: int, n_splits: int, embargo: int) -> list[list[int]]
Example 1
Inputwalk_forward_folds(10, 2, 0)
Output[[0, 4, 4, 7], [0, 7, 7, 10]]
Example 2
Inputwalk_forward_folds(10, 2, 2)
Output[[0, 2, 4, 7], [0, 5, 7, 10]]