Fix the walk-forward split for a trading backtest
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_rowsorn_splitsis 0 or less, or ifembargois negative. - The block size is
n_rowsdivided byn_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) * sizeand ends itsizerows later, so the last block ends atn_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
Input
walk_forward_folds(10, 2, 0)Output
[[0, 4, 4, 7], [0, 7, 7, 10]]Example 2
Input
walk_forward_folds(10, 2, 2)Output
[[0, 2, 4, 7], [0, 5, 7, 10]]