Skip to content

Commit

Permalink
Uses last bar in test split when windows=1
Browse files Browse the repository at this point in the history
Guarantees that the last bar of data is used in the test split when the number of train/test windows is 1.
  • Loading branch information
edtechre committed Aug 12, 2024
1 parent 90c9f76 commit e123add
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 4 deletions.
8 changes: 5 additions & 3 deletions src/pybroker/strategy.py
Original file line number Diff line number Diff line change
Expand Up @@ -698,12 +698,14 @@ def walkforward_split(
raise ValueError(error_msg)
train_length = int(res * train_size)
test_length = int(res * (1 - train_size))
train_start = 0
train_end = train_length
train_start = (
len(window_dates) - lookahead - train_length - test_length - 1
)
train_end = train_start + train_length
test_start = train_end + lookahead
if test_start >= len(window_dates):
raise ValueError(error_msg)
test_end = test_start + test_length
test_end = len(window_dates) - 1
train_idx = dates[
(dates[date_col] >= window_dates[train_start])
& (dates[date_col] <= window_dates[train_end])
Expand Down
4 changes: 3 additions & 1 deletion tests/test_strategy.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ def _verify_windows(
)
dates = sorted(dates)
assert len(results) == windows
for train_idx, test_idx in results:
for i, (train_idx, test_idx) in enumerate(results):
assert len(dates) - (len(train_idx) + len(test_idx) * windows) >= 0
assert not (set(train_idx) & set(test_idx))
assert len(train_idx) or len(test_idx)
Expand All @@ -112,6 +112,8 @@ def _verify_windows(
assert dates[train_end_index - 2] != dates[test_start_index]
if train_size == 0.5:
assert len(train_idx) == len(test_idx)
if len(test_idx) and i == len(results) - 1:
assert dates[dates_length - 1] == dates[sorted(test_idx)[-1]]

@pytest.mark.parametrize(
"dates_length, windows, lookahead, train_size",
Expand Down

0 comments on commit e123add

Please sign in to comment.